retry.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213
  1. # Copyright 2016–2021 Julien Danjou
  2. # Copyright 2016 Joshua Harlow
  3. # Copyright 2013-2014 Ray Holder
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. import abc
  17. import re
  18. import typing
  19. if typing.TYPE_CHECKING:
  20. from pip._vendor.tenacity import RetryCallState
  21. class retry_base(abc.ABC):
  22. """Abstract base class for retry strategies."""
  23. @abc.abstractmethod
  24. def __call__(self, retry_state: "RetryCallState") -> bool:
  25. pass
  26. def __and__(self, other: "retry_base") -> "retry_all":
  27. return retry_all(self, other)
  28. def __or__(self, other: "retry_base") -> "retry_any":
  29. return retry_any(self, other)
  30. class _retry_never(retry_base):
  31. """Retry strategy that never rejects any result."""
  32. def __call__(self, retry_state: "RetryCallState") -> bool:
  33. return False
  34. retry_never = _retry_never()
  35. class _retry_always(retry_base):
  36. """Retry strategy that always rejects any result."""
  37. def __call__(self, retry_state: "RetryCallState") -> bool:
  38. return True
  39. retry_always = _retry_always()
  40. class retry_if_exception(retry_base):
  41. """Retry strategy that retries if an exception verifies a predicate."""
  42. def __init__(self, predicate: typing.Callable[[BaseException], bool]) -> None:
  43. self.predicate = predicate
  44. def __call__(self, retry_state: "RetryCallState") -> bool:
  45. if retry_state.outcome.failed:
  46. return self.predicate(retry_state.outcome.exception())
  47. else:
  48. return False
  49. class retry_if_exception_type(retry_if_exception):
  50. """Retries if an exception has been raised of one or more types."""
  51. def __init__(
  52. self,
  53. exception_types: typing.Union[
  54. typing.Type[BaseException],
  55. typing.Tuple[typing.Type[BaseException], ...],
  56. ] = Exception,
  57. ) -> None:
  58. self.exception_types = exception_types
  59. super().__init__(lambda e: isinstance(e, exception_types))
  60. class retry_if_not_exception_type(retry_if_exception):
  61. """Retries except an exception has been raised of one or more types."""
  62. def __init__(
  63. self,
  64. exception_types: typing.Union[
  65. typing.Type[BaseException],
  66. typing.Tuple[typing.Type[BaseException], ...],
  67. ] = Exception,
  68. ) -> None:
  69. self.exception_types = exception_types
  70. super().__init__(lambda e: not isinstance(e, exception_types))
  71. class retry_unless_exception_type(retry_if_exception):
  72. """Retries until an exception is raised of one or more types."""
  73. def __init__(
  74. self,
  75. exception_types: typing.Union[
  76. typing.Type[BaseException],
  77. typing.Tuple[typing.Type[BaseException], ...],
  78. ] = Exception,
  79. ) -> None:
  80. self.exception_types = exception_types
  81. super().__init__(lambda e: not isinstance(e, exception_types))
  82. def __call__(self, retry_state: "RetryCallState") -> bool:
  83. # always retry if no exception was raised
  84. if not retry_state.outcome.failed:
  85. return True
  86. return self.predicate(retry_state.outcome.exception())
  87. class retry_if_result(retry_base):
  88. """Retries if the result verifies a predicate."""
  89. def __init__(self, predicate: typing.Callable[[typing.Any], bool]) -> None:
  90. self.predicate = predicate
  91. def __call__(self, retry_state: "RetryCallState") -> bool:
  92. if not retry_state.outcome.failed:
  93. return self.predicate(retry_state.outcome.result())
  94. else:
  95. return False
  96. class retry_if_not_result(retry_base):
  97. """Retries if the result refutes a predicate."""
  98. def __init__(self, predicate: typing.Callable[[typing.Any], bool]) -> None:
  99. self.predicate = predicate
  100. def __call__(self, retry_state: "RetryCallState") -> bool:
  101. if not retry_state.outcome.failed:
  102. return not self.predicate(retry_state.outcome.result())
  103. else:
  104. return False
  105. class retry_if_exception_message(retry_if_exception):
  106. """Retries if an exception message equals or matches."""
  107. def __init__(
  108. self,
  109. message: typing.Optional[str] = None,
  110. match: typing.Optional[str] = None,
  111. ) -> None:
  112. if message and match:
  113. raise TypeError(f"{self.__class__.__name__}() takes either 'message' or 'match', not both")
  114. # set predicate
  115. if message:
  116. def message_fnc(exception: BaseException) -> bool:
  117. return message == str(exception)
  118. predicate = message_fnc
  119. elif match:
  120. prog = re.compile(match)
  121. def match_fnc(exception: BaseException) -> bool:
  122. return bool(prog.match(str(exception)))
  123. predicate = match_fnc
  124. else:
  125. raise TypeError(f"{self.__class__.__name__}() missing 1 required argument 'message' or 'match'")
  126. super().__init__(predicate)
  127. class retry_if_not_exception_message(retry_if_exception_message):
  128. """Retries until an exception message equals or matches."""
  129. def __init__(
  130. self,
  131. message: typing.Optional[str] = None,
  132. match: typing.Optional[str] = None,
  133. ) -> None:
  134. super().__init__(message, match)
  135. # invert predicate
  136. if_predicate = self.predicate
  137. self.predicate = lambda *args_, **kwargs_: not if_predicate(*args_, **kwargs_)
  138. def __call__(self, retry_state: "RetryCallState") -> bool:
  139. if not retry_state.outcome.failed:
  140. return True
  141. return self.predicate(retry_state.outcome.exception())
  142. class retry_any(retry_base):
  143. """Retries if any of the retries condition is valid."""
  144. def __init__(self, *retries: retry_base) -> None:
  145. self.retries = retries
  146. def __call__(self, retry_state: "RetryCallState") -> bool:
  147. return any(r(retry_state) for r in self.retries)
  148. class retry_all(retry_base):
  149. """Retries if all the retries condition are valid."""
  150. def __init__(self, *retries: retry_base) -> None:
  151. self.retries = retries
  152. def __call__(self, retry_state: "RetryCallState") -> bool:
  153. return all(r(retry_state) for r in self.retries)