From 88bc6a46a6ab03ccfa1aab570cf45f5ff9889d44 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 6 Apr 2023 13:58:26 +0200
Subject: [PATCH] Store self user in entity cache

---
 telethon/_updates/entitycache.py      |  4 +++-
 telethon/client/auth.py               |  6 ++----
 telethon/client/telegrambaseclient.py |  4 ----
 telethon/client/updates.py            |  4 ++--
 telethon/client/users.py              | 21 +++++++++------------
 5 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/telethon/_updates/entitycache.py b/telethon/_updates/entitycache.py
index 0c28d3a2..3a8f49ed 100644
--- a/telethon/_updates/entitycache.py
+++ b/telethon/_updates/entitycache.py
@@ -15,9 +15,11 @@ class EntityCache:
         self.self_id = self_id
         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_bot = bot
+        if hash:
+            self.hash_map[id] = (hash, EntityType.BOT if bot else EntityType.USER)
 
     def get(self, id):
         try:
diff --git a/telethon/client/auth.py b/telethon/client/auth.py
index 4bf20dbb..3cdf998e 100644
--- a/telethon/client/auth.py
+++ b/telethon/client/auth.py
@@ -380,8 +380,7 @@ class AuthMethods:
 
         Returns the input user parameter.
         """
-        self._bot = bool(user.bot)
-        self._self_input_peer = utils.get_input_peer(user, allow_self=False)
+        self._mb_entity_cache.set_self_user(user.id, user.bot, user.access_hash)
         self._authorized = True
 
         state = await self(functions.updates.GetStateRequest())
@@ -531,8 +530,7 @@ class AuthMethods:
         except errors.RPCError:
             return False
 
-        self._bot = None
-        self._self_input_peer = None
+        self._mb_entity_cache.set_self_user(None, False, None)
         self._authorized = False
 
         await self.disconnect()
diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py
index 63d69bb8..4d0bb800 100644
--- a/telethon/client/telegrambaseclient.py
+++ b/telethon/client/telegrambaseclient.py
@@ -435,10 +435,6 @@ class TelegramBaseClient(abc.ABC):
         self._phone = 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`)
         self._megagroup_cache = {}
 
diff --git a/telethon/client/updates.py b/telethon/client/updates.py
index 6e659ec3..a8ad3a10 100644
--- a/telethon/client/updates.py
+++ b/telethon/client/updates.py
@@ -495,10 +495,10 @@ class UpdateMethods:
         # TODO only used for AlbumHack, and MessageBox is not really designed for this
         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
             # 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
             # fine, we will just retry next time anyway.
diff --git a/telethon/client/users.py b/telethon/client/users.py
index 50f9ff81..bbb264fb 100644
--- a/telethon/client/users.py
+++ b/telethon/client/users.py
@@ -152,20 +152,17 @@ class UserMethods:
                 me = await client.get_me()
                 print(me.username)
         """
-        if input_peer and self._self_input_peer:
-            return self._self_input_peer
+        if input_peer and self._mb_entity_cache.self_id:
+            return self._mb_entity_cache.get(self._mb_entity_cache.self_id)._as_input_peer()
 
         try:
             me = (await self(
                 functions.users.GetUsersRequest([types.InputUserSelf()])))[0]
 
-            self._bot = me.bot
-            if not self._self_input_peer:
-                self._self_input_peer = utils.get_input_peer(
-                    me, allow_self=False
-                )
+            if not self._mb_entity_cache.self_id:
+                self._mb_entity_cache.set_self_user(me.id, me.bot, me.access_hash)
 
-            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:
             return None
 
@@ -177,7 +174,7 @@ class UserMethods:
         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.
         """
-        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:
         """
@@ -191,10 +188,10 @@ class UserMethods:
                 else:
                     print('Hello')
         """
-        if self._bot is None:
-            self._bot = (await self.get_me()).bot
+        if self._mb_entity_cache.self_id is None:
+            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:
         """