Review, unify and simplify retry_range usage

This commit is contained in:
Lonami Exo 2022-01-24 11:29:02 +01:00
parent b0b1f30436
commit a25f019964
2 changed files with 7 additions and 20 deletions

View File

@ -102,23 +102,11 @@ def strip_text(text, entities):
return text return text
def retry_range(retries, force_retry=True): def retry_range(retries):
""" """
Generates an integer sequence starting from 1. If `retries` is Generates an integer sequence starting from 1, always returning once, and adding the given retries.
not a zero or a positive integer value, the sequence will be
infinite, otherwise it will end at `retries + 1`.
""" """
return range(1, max(retries, 0) + 2)
# We need at least one iteration even if the retries are 0
# when force_retry is True.
if force_retry and not (retries is None or retries < 0):
retries += 1
attempt = 0
while attempt != retries:
attempt += 1
yield attempt
async def _maybe_await(value): async def _maybe_await(value):

View File

@ -249,9 +249,9 @@ class MTProtoSender:
break # all steps done, break retry loop break # all steps done, break retry loop
else: else:
if not connected: if not connected:
raise ConnectionError('Connection to Telegram failed {} time(s)'.format(self._retries)) raise ConnectionError('Connection to Telegram failed {} time(s)'.format(1 + self._retries))
e = ConnectionError('auth_key generation failed {} time(s)'.format(self._retries)) e = ConnectionError('auth_key generation failed {} time(s)'.format(1 + self._retries))
await self._disconnect(error=e) await self._disconnect(error=e)
raise e raise e
@ -349,12 +349,11 @@ class MTProtoSender:
# Start with a clean state (and thus session ID) to avoid old msgs # Start with a clean state (and thus session ID) to avoid old msgs
self._state.reset() self._state.reset()
retries = self._retries if self._auto_reconnect else 0 retry_range = helpers.retry_range(self._retries) if self._auto_reconnect else range(0)
attempt = 0 attempt = 0
ok = True ok = True
# We're already "retrying" to connect, so we don't want to force retries for attempt in retry_range:
for attempt in helpers.retry_range(retries, force_retry=False):
try: try:
await self._connect() await self._connect()
except (IOError, asyncio.TimeoutError) as e: except (IOError, asyncio.TimeoutError) as e: