Store self user in entity cache

This commit is contained in:
Lonami Exo 2023-04-06 13:58:26 +02:00
parent 97b0ba6707
commit 88bc6a46a6
5 changed files with 16 additions and 23 deletions

View File

@ -15,9 +15,11 @@ class EntityCache:
self.self_id = self_id self.self_id = self_id
self.self_bot = self_bot self.self_bot = self_bot
def set_self_user(self, id, bot): def set_self_user(self, id, bot, hash):
self.self_id = id self.self_id = id
self.self_bot = bot self.self_bot = bot
if hash:
self.hash_map[id] = (hash, EntityType.BOT if bot else EntityType.USER)
def get(self, id): def get(self, id):
try: try:

View File

@ -380,8 +380,7 @@ class AuthMethods:
Returns the input user parameter. Returns the input user parameter.
""" """
self._bot = bool(user.bot) self._mb_entity_cache.set_self_user(user.id, user.bot, user.access_hash)
self._self_input_peer = utils.get_input_peer(user, allow_self=False)
self._authorized = True self._authorized = True
state = await self(functions.updates.GetStateRequest()) state = await self(functions.updates.GetStateRequest())
@ -531,8 +530,7 @@ class AuthMethods:
except errors.RPCError: except errors.RPCError:
return False return False
self._bot = None self._mb_entity_cache.set_self_user(None, False, None)
self._self_input_peer = None
self._authorized = False self._authorized = False
await self.disconnect() await self.disconnect()

View File

@ -435,10 +435,6 @@ class TelegramBaseClient(abc.ABC):
self._phone = None self._phone = None
self._tos = None self._tos = None
# Sometimes we need to know who we are, cache the self peer
self._self_input_peer = None
self._bot = None
# A place to store if channels are a megagroup or not (see `edit_admin`) # A place to store if channels are a megagroup or not (see `edit_admin`)
self._megagroup_cache = {} self._megagroup_cache = {}

View File

@ -495,10 +495,10 @@ class UpdateMethods:
# TODO only used for AlbumHack, and MessageBox is not really designed for this # TODO only used for AlbumHack, and MessageBox is not really designed for this
others = None others = None
if not self._self_input_peer: if not self._mb_entity_cache.self_id:
# Some updates require our own ID, so we must make sure # Some updates require our own ID, so we must make sure
# that the event builder has offline access to it. Calling # that the event builder has offline access to it. Calling
# `get_me()` will cache it under `self._self_input_peer`. # `get_me()` will cache it under `self._mb_entity_cache`.
# #
# It will return `None` if we haven't logged in yet which is # It will return `None` if we haven't logged in yet which is
# fine, we will just retry next time anyway. # fine, we will just retry next time anyway.

View File

@ -152,20 +152,17 @@ class UserMethods:
me = await client.get_me() me = await client.get_me()
print(me.username) print(me.username)
""" """
if input_peer and self._self_input_peer: if input_peer and self._mb_entity_cache.self_id:
return self._self_input_peer return self._mb_entity_cache.get(self._mb_entity_cache.self_id)._as_input_peer()
try: try:
me = (await self( me = (await self(
functions.users.GetUsersRequest([types.InputUserSelf()])))[0] functions.users.GetUsersRequest([types.InputUserSelf()])))[0]
self._bot = me.bot if not self._mb_entity_cache.self_id:
if not self._self_input_peer: self._mb_entity_cache.set_self_user(me.id, me.bot, me.access_hash)
self._self_input_peer = utils.get_input_peer(
me, allow_self=False
)
return self._self_input_peer if input_peer else me return utils.get_input_peer(me, allow_self=False) if input_peer else me
except errors.UnauthorizedError: except errors.UnauthorizedError:
return None return None
@ -177,7 +174,7 @@ class UserMethods:
This property is used in every update, and some like `updateLoginToken` This property is used in every update, and some like `updateLoginToken`
occur prior to login, so it gracefully handles when no ID is known yet. occur prior to login, so it gracefully handles when no ID is known yet.
""" """
return self._self_input_peer.user_id if self._self_input_peer else None return self._mb_entity_cache.self_id
async def is_bot(self: 'TelegramClient') -> bool: async def is_bot(self: 'TelegramClient') -> bool:
""" """
@ -191,10 +188,10 @@ class UserMethods:
else: else:
print('Hello') print('Hello')
""" """
if self._bot is None: if self._mb_entity_cache.self_id is None:
self._bot = (await self.get_me()).bot await self.get_me(input_peer=True)
return self._bot return self._mb_entity_cache.self_bot
async def is_user_authorized(self: 'TelegramClient') -> bool: async def is_user_authorized(self: 'TelegramClient') -> bool:
""" """