_asyncio.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. # -*- coding: utf-8 -*-
  2. # Copyright 2016 Étienne Bersac
  3. # Copyright 2016 Julien Danjou
  4. # Copyright 2016 Joshua Harlow
  5. # Copyright 2013-2014 Ray Holder
  6. #
  7. # Licensed under the Apache License, Version 2.0 (the "License");
  8. # you may not use this file except in compliance with the License.
  9. # You may obtain a copy of the License at
  10. #
  11. # http://www.apache.org/licenses/LICENSE-2.0
  12. #
  13. # Unless required by applicable law or agreed to in writing, software
  14. # distributed under the License is distributed on an "AS IS" BASIS,
  15. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16. # See the License for the specific language governing permissions and
  17. # limitations under the License.
  18. import sys
  19. from asyncio import sleep
  20. from pip._vendor.tenacity import AttemptManager
  21. from pip._vendor.tenacity import BaseRetrying
  22. from pip._vendor.tenacity import DoAttempt
  23. from pip._vendor.tenacity import DoSleep
  24. from pip._vendor.tenacity import RetryCallState
  25. class AsyncRetrying(BaseRetrying):
  26. def __init__(self, sleep=sleep, **kwargs):
  27. super(AsyncRetrying, self).__init__(**kwargs)
  28. self.sleep = sleep
  29. async def __call__(self, fn, *args, **kwargs):
  30. self.begin(fn)
  31. retry_state = RetryCallState(retry_object=self, fn=fn, args=args, kwargs=kwargs)
  32. while True:
  33. do = self.iter(retry_state=retry_state)
  34. if isinstance(do, DoAttempt):
  35. try:
  36. result = await fn(*args, **kwargs)
  37. except BaseException: # noqa: B902
  38. retry_state.set_exception(sys.exc_info())
  39. else:
  40. retry_state.set_result(result)
  41. elif isinstance(do, DoSleep):
  42. retry_state.prepare_for_next_attempt()
  43. await self.sleep(do)
  44. else:
  45. return do
  46. def __aiter__(self):
  47. self.begin(None)
  48. self._retry_state = RetryCallState(self, fn=None, args=(), kwargs={})
  49. return self
  50. async def __anext__(self):
  51. while True:
  52. do = self.iter(retry_state=self._retry_state)
  53. if do is None:
  54. raise StopAsyncIteration
  55. elif isinstance(do, DoAttempt):
  56. return AttemptManager(retry_state=self._retry_state)
  57. elif isinstance(do, DoSleep):
  58. self._retry_state.prepare_for_next_attempt()
  59. await self.sleep(do)
  60. else:
  61. return do
  62. def wraps(self, fn):
  63. fn = super().wraps(fn)
  64. # Ensure wrapper is recognized as a coroutine function.
  65. async def async_wrapped(*args, **kwargs):
  66. return await fn(*args, **kwargs)
  67. # Preserve attributes
  68. async_wrapped.retry = fn.retry
  69. async_wrapped.retry_with = fn.retry_with
  70. return async_wrapped