Use the new in-memory entity cache

This should avoid a disk access every time an input entity
is needed, which is very often. Another step for #1141.

telethon/client/telegrambaseclient.py

@@ -235,7 +235,6 @@ class TelegramBaseClient(abc.ABC):
         #
         # The session files only wants the entities to persist
         # them to disk, and to save additional useful information.
-        # TODO Make use of _entity_cache
         # TODO Session should probably return all cached
         #      info of entities, not just the input versions
         self.session = session

telethon/client/users.py

@@ -308,8 +308,8 @@ class UserMethods(TelegramBaseClient):
 
         try:
             # First try to get the entity from cache, otherwise figure it out
-            return self.session.get_input_entity(peer)
-        except ValueError:
+            return self._entity_cache[peer]
+        except KeyError:
             pass
 
         if isinstance(peer, str):

telethon/entitycache.py

@@ -31,15 +31,17 @@ class EntityCache:
     def __getitem__(self, item):
         """
         Gets the corresponding :tl:`InputPeer` for the given ID or peer,
-        or returns `None` on error/not found.
+        or raises ``KeyError`` on any error (i.e. cannot be found).
         """
         if not isinstance(item, int) or item < 0:
             try:
-                return self.__dict__.get(utils.get_peer_id(item))
+                return self.__dict__[utils.get_peer_id(item)]
             except TypeError:
-                return None
+                raise KeyError('Invalid key will not have entity') from None
 
         for cls in (types.PeerUser, types.PeerChat, types.PeerChannel):
-            result = self.__dict__.get(cls(item))
+            result = self.__dict__.get(utils.get_peer_id(cls(item)))
             if result:
                 return result
+
+        raise KeyError('No cached entity for the given key')

telethon/events/callbackquery.py

@@ -156,10 +156,8 @@ class CallbackQuery(EventBuilder):
             if not getattr(self._input_sender, 'access_hash', True):
                 # getattr with True to handle the InputPeerSelf() case
                 try:
-                    self._input_sender = self._client.session.get_input_entity(
-                        self._sender_id
-                    )
-                except ValueError:
+                    self._input_sender = self._client._entity_cache[self._sender_id]
+                except KeyError:
                     m = await self.get_message()
                     if m:
                         self._sender = m._sender

telethon/events/chataction.py

@@ -362,10 +362,8 @@ class ChatAction(EventBuilder):
                 self._input_users = []
                 for peer in self._user_peers:
                     try:
-                        self._input_users.append(
-                            self._client.session.get_input_entity(peer)
-                        )
-                    except ValueError:
+                        self._input_users.append(self._client._entity_cache[peer])
+                    except KeyError:
                         pass
             return self._input_users or []
 

telethon/events/common.py

@@ -154,10 +154,8 @@ class EventCommon(ChatGetter, abc.ABC):
             self._input_chat = utils.get_input_peer(self._chat)
         except TypeError:
             try:
-                self._input_chat = self._client.session.get_input_entity(
-                    self._chat_peer
-                )
-            except ValueError:
+                self._input_chat = self._client._entity_cache[self._chat_peer]
+            except KeyError:
                 self._input_chat = None
 
     @property

telethon/events/userupdate.py

@@ -170,7 +170,7 @@ class UserUpdate(EventBuilder):
         def _set_client(self, client):
             if isinstance(self._chat_peer, int):
                 try:
-                    chat = client.session.get_input_entity(self._chat_peer)
+                    chat = client._entity_cache[self._chat_peer]
                     if isinstance(chat, types.InputPeerChat):
                         self._chat_peer = types.PeerChat(self._chat_peer)
                     elif isinstance(chat, types.InputPeerChannel):
@@ -178,7 +178,7 @@ class UserUpdate(EventBuilder):
                     else:
                         # Should not happen
                         self._chat_peer = types.PeerUser(self._chat_peer)
-                except ValueError:
+                except KeyError:
                     # Hope for the best. We don't know where this event
                     # occurred but it was most likely in a channel.
                     self._chat_peer = types.PeerChannel(self._chat_peer)

telethon/tl/custom/chatgetter.py

@@ -49,9 +49,8 @@ class ChatGetter(abc.ABC):
         """
         if self._input_chat is None and self._chat_peer:
             try:
-                self._input_chat =\
-                    self._client.session.get_input_entity(self._chat_peer)
-            except ValueError:
+                self._input_chat = self._client._entity_cache(self._chat_peer)
+            except KeyError:
                 pass
 
         return self._input_chat

telethon/tl/custom/draft.py

@@ -64,9 +64,8 @@ class Draft:
         """
         if not self._input_entity:
             try:
-                self._input_entity =\
-                    self._client.session.get_input_entity(self._peer)
-            except ValueError:
+                self._input_entity = self._client._entity_cache[self._peer]
+            except KeyError:
                 pass
 
         return self._input_entity

telethon/tl/custom/message.py

@@ -834,8 +834,10 @@ class Message(ChatGetter, SenderGetter, TLObject, abc.ABC):
                         if not bot:
                             raise ValueError('No input sender')
                     else:
-                        return self._client.session.get_input_entity(
-                            self.via_bot_id)
+                        try:
+                            return self._client._entity_cache[self.via_bot_id]
+                        except KeyError:
+                            raise ValueError('No input sender') from None
 
     def _document_by_attribute(self, kind, condition=None):
         """