Remove redundant entity cache

Progress towards #3989.
May also help with #3235.

telethon/tl/custom/chatgetter.py

@@ -66,8 +66,9 @@ class ChatGetter(abc.ABC):
         """
         if self._input_chat is None and self._chat_peer and self._client:
             try:
-                self._input_chat = self._client._entity_cache[self._chat_peer]
-            except KeyError:
+                self._input_chat = self._client._mb_entity_cache.get(
+                        utils.get_peer_id(self._chat_peer, add_mark=False))._as_input_peer()
+            except AttributeError:
                 pass
 
         return self._input_chat

telethon/tl/custom/draft.py

@@ -5,7 +5,7 @@ from ..functions.messages import SaveDraftRequest
 from ..types import DraftMessage
 from ...errors import RPCError
 from ...extensions import markdown
-from ...utils import get_input_peer, get_peer
+from ...utils import get_input_peer, get_peer, get_peer_id
 
 
 class Draft:
@@ -53,8 +53,9 @@ class Draft:
         """
         if not self._input_entity:
             try:
-                self._input_entity = self._client._entity_cache[self._peer]
-            except KeyError:
+                self._input_entity = self._client._mb_entity_cache.get(
+                        get_peer_id(self._peer, add_mark=False))._as_input_peer()
+            except AttributeError:
                 pass
 
         return self._input_entity

telethon/tl/custom/forward.py

@@ -36,12 +36,12 @@ class Forward(ChatGetter, SenderGetter):
             if ty == helpers._EntityType.USER:
                 sender_id = utils.get_peer_id(original.from_id)
                 sender, input_sender = utils._get_entity_pair(
-                    sender_id, entities, client._entity_cache)
+                    sender_id, entities, client._mb_entity_cache)
 
             elif ty in (helpers._EntityType.CHAT, helpers._EntityType.CHANNEL):
                 peer = original.from_id
                 chat, input_chat = utils._get_entity_pair(
-                    utils.get_peer_id(peer), entities, client._entity_cache)
+                    utils.get_peer_id(peer), entities, client._mb_entity_cache)
 
         # This call resets the client
         ChatGetter.__init__(self, peer, chat=chat, input_chat=input_chat)

telethon/tl/custom/message.py

@@ -285,7 +285,7 @@ class Message(ChatGetter, SenderGetter, TLObject):
         if self.peer_id == types.PeerUser(client._self_id) and not self.fwd_from:
             self.out = True
 
-        cache = client._entity_cache
+        cache = client._mb_entity_cache
 
         self._sender, self._input_sender = utils._get_entity_pair(
             self.sender_id, entities, cache)
@@ -1138,8 +1138,9 @@ class Message(ChatGetter, SenderGetter, TLObject):
                         return bot
                     else:
                         try:
-                            return self._client._entity_cache[self.via_bot_id]
-                        except KeyError:
+                            return self._client._mb_entity_cache.get(
+                                utils.resolve_id(self.via_bot_id)[0])._as_input_peer()
+                        except AttributeError:
                             raise ValueError('No input sender') from None
 
     def _document_by_attribute(self, kind, condition=None):

telethon/tl/custom/sendergetter.py

@@ -1,5 +1,7 @@
 import abc
 
+from ... import utils
+
 
 class SenderGetter(abc.ABC):
     """
@@ -69,9 +71,9 @@ class SenderGetter(abc.ABC):
         """
         if self._input_sender is None and self._sender_id and self._client:
             try:
-                self._input_sender = \
-                    self._client._entity_cache[self._sender_id]
-            except KeyError:
+                self._input_sender = self._client._mb_entity_cache.get(
+                        utils.resolve_id(self._sender_id)[0])._as_input_peer()
+            except AttributeError:
                 pass
         return self._input_sender