Raise on phone/network migrate if authorized

This commit is contained in:
Lonami Exo 2018-06-14 23:13:42 +02:00
parent c788e17c22
commit 06cdf1fdb0
2 changed files with 14 additions and 3 deletions

View File

@ -159,9 +159,6 @@ class AuthMethods(MessageParseMethods, UserMethods):
return self return self
async def is_user_authorized(self):
return await self.get_me() is not None
async def sign_in( async def sign_in(
self, phone=None, code=None, password=None, self, phone=None, code=None, password=None,
bot_token=None, phone_code_hash=None): bot_token=None, phone_code_hash=None):
@ -239,6 +236,7 @@ class AuthMethods(MessageParseMethods, UserMethods):
self._self_input_peer = utils.get_input_peer( self._self_input_peer = utils.get_input_peer(
result.user, allow_self=False result.user, allow_self=False
) )
self._authorized = True
return result.user return result.user
async def sign_up(self, code, first_name, last_name=''): async def sign_up(self, code, first_name, last_name=''):
@ -291,6 +289,7 @@ class AuthMethods(MessageParseMethods, UserMethods):
self._self_input_peer = utils.get_input_peer( self._self_input_peer = utils.get_input_peer(
result.user, allow_self=False result.user, allow_self=False
) )
self._authorized = True
return result.user return result.user
async def send_code_request(self, phone, force_sms=False): async def send_code_request(self, phone, force_sms=False):

View File

@ -35,6 +35,10 @@ class UserMethods(TelegramBaseClient):
raise raise
except (errors.PhoneMigrateError, errors.NetworkMigrateError, except (errors.PhoneMigrateError, errors.NetworkMigrateError,
errors.UserMigrateError) as e: errors.UserMigrateError) as e:
should_raise = isinstance(errors.PhoneMigrateError,
errors.NetworkMigrateError)
if should_raise and await self.is_user_authorized():
raise
await self._switch_dc(e.new_dc) await self._switch_dc(e.new_dc)
raise ValueError('Number of retries reached 0') raise ValueError('Number of retries reached 0')
@ -71,6 +75,14 @@ class UserMethods(TelegramBaseClient):
except errors.UnauthorizedError: except errors.UnauthorizedError:
return None return None
async def is_user_authorized(self):
"""
Returns ``True`` if the user is authorized.
"""
if self._authorized is None:
self._authorized = await self.get_me() is not None
return self._authorized
async def get_entity(self, entity): async def get_entity(self, entity):
""" """
Turns the given entity into a valid Telegram user or chat. Turns the given entity into a valid Telegram user or chat.