123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081 |
- # -*- coding: utf-8 -*-
- # Copyright 2016 Étienne Bersac
- # Copyright 2016 Julien Danjou
- # Copyright 2016 Joshua Harlow
- # Copyright 2013-2014 Ray Holder
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- import sys
- from asyncio import sleep
- from pip._vendor.tenacity import AttemptManager
- from pip._vendor.tenacity import BaseRetrying
- from pip._vendor.tenacity import DoAttempt
- from pip._vendor.tenacity import DoSleep
- from pip._vendor.tenacity import RetryCallState
- class AsyncRetrying(BaseRetrying):
- def __init__(self, sleep=sleep, **kwargs):
- super(AsyncRetrying, self).__init__(**kwargs)
- self.sleep = sleep
- async def __call__(self, fn, *args, **kwargs):
- self.begin(fn)
- retry_state = RetryCallState(retry_object=self, fn=fn, args=args, kwargs=kwargs)
- while True:
- do = self.iter(retry_state=retry_state)
- if isinstance(do, DoAttempt):
- try:
- result = await fn(*args, **kwargs)
- except BaseException: # noqa: B902
- retry_state.set_exception(sys.exc_info())
- else:
- retry_state.set_result(result)
- elif isinstance(do, DoSleep):
- retry_state.prepare_for_next_attempt()
- await self.sleep(do)
- else:
- return do
- def __aiter__(self):
- self.begin(None)
- self._retry_state = RetryCallState(self, fn=None, args=(), kwargs={})
- return self
- async def __anext__(self):
- while True:
- do = self.iter(retry_state=self._retry_state)
- if do is None:
- raise StopAsyncIteration
- elif isinstance(do, DoAttempt):
- return AttemptManager(retry_state=self._retry_state)
- elif isinstance(do, DoSleep):
- self._retry_state.prepare_for_next_attempt()
- await self.sleep(do)
- else:
- return do
- def wraps(self, fn):
- fn = super().wraps(fn)
- # Ensure wrapper is recognized as a coroutine function.
- async def async_wrapped(*args, **kwargs):
- return await fn(*args, **kwargs)
- # Preserve attributes
- async_wrapped.retry = fn.retry
- async_wrapped.retry_with = fn.retry_with
- return async_wrapped
|