Support getting any entity by just their positive ID

telethon/sessions/memory.py

@@ -1,9 +1,8 @@
 from enum import Enum
 
-from .. import utils
 from .abstract import Session
+from .. import utils
 from ..tl import TLObject
-
 from ..tl.types import (
     PeerUser, PeerChat, PeerChannel,
     InputPeerUser, InputPeerChat, InputPeerChannel,
@@ -148,10 +147,19 @@ class MemorySession(Session):
         except StopIteration:
             pass
 
-    def get_entity_rows_by_id(self, id):
+    def get_entity_rows_by_id(self, id, exact=True):
         try:
-            return next((id, hash) for found_id, hash, _, _, _
-                        in self._entities if found_id == id)
+            if exact:
+                return next((id, hash) for found_id, hash, _, _, _
+                            in self._entities if found_id == id)
+            else:
+                ids = (
+                    utils.get_peer_id(PeerUser(id)),
+                    utils.get_peer_id(PeerChat(id)),
+                    utils.get_peer_id(PeerChannel(id))
+                )
+                return next((id, hash) for found_id, hash, _, _, _
+                            in self._entities if found_id in ids)
         except StopIteration:
             pass
 
@@ -167,6 +175,9 @@ class MemorySession(Session):
             # Not a TLObject or can't be cast into InputPeer
             if isinstance(key, TLObject):
                 key = utils.get_peer_id(key)
+                exact = True
+            else:
+                exact = False
 
         result = None
         if isinstance(key, str):
@@ -178,8 +189,8 @@ class MemorySession(Session):
                 if username:
                     result = self.get_entity_rows_by_username(username)
 
-        if isinstance(key, int):
-            result = self.get_entity_rows_by_id(key)
+        elif isinstance(key, int):
+            result = self.get_entity_rows_by_id(key, exact)
 
         if not result and isinstance(key, str):
             result = self.get_entity_rows_by_name(key)

telethon/sessions/sqlalchemy.py

@@ -4,12 +4,13 @@ try:
     import sqlalchemy as sql
 except ImportError:
     sql = None
-    pass
-
-from ..crypto import AuthKey
-from ..tl.types import InputPhoto, InputDocument
 
 from .memory import MemorySession, _SentFileType
+from .. import utils
+from ..crypto import AuthKey
+from ..tl.types import (
+    InputPhoto, InputDocument, PeerUser, PeerChat, PeerChannel
+)
 
 LATEST_VERSION = 1
 
@@ -201,8 +202,18 @@ class AlchemySession(MemorySession):
                              self.Entity.name == key).one_or_none()
         return row.id, row.hash if row else None
 
-    def get_entity_rows_by_id(self, key):
-        row = self._db_query(self.Entity, self.Entity.id == key).one_or_none()
+    def get_entity_rows_by_id(self, key, exact=True):
+        if exact:
+            query = self._db_query(self.Entity, self.Entity.id == key)
+        else:
+            ids = (
+                utils.get_peer_id(PeerUser(key)),
+                utils.get_peer_id(PeerChat(key)),
+                utils.get_peer_id(PeerChannel(key))
+            )
+            query = self._db_query(self.Entity, self.Entity.id in ids)
+
+        row = query.one_or_none()
         return row.id, row.hash if row else None
 
     def get_file(self, md5_digest, file_size, cls):

telethon/sessions/sqlite.py

@@ -6,9 +6,10 @@ from os.path import isfile as file_exists
 from threading import Lock, RLock
 
 from .memory import MemorySession, _SentFileType
+from .. import utils
 from ..crypto import AuthKey
 from ..tl.types import (
-    InputPhoto, InputDocument
+    InputPhoto, InputDocument, PeerUser, PeerChat, PeerChannel
 )
 
 EXTENSION = '.session'
@@ -282,9 +283,19 @@ class SQLiteSession(MemorySession):
         return self._fetchone_entity(
             'select id, hash from entities where name=?', (name,))
 
-    def get_entity_rows_by_id(self, id):
-        return self._fetchone_entity(
-            'select id, hash from entities where id=?', (id,))
+    def get_entity_rows_by_id(self, id, exact=True):
+        if exact:
+            return self._fetchone_entity(
+                'select id, hash from entities where id=?', (id,))
+        else:
+            ids = (
+                utils.get_peer_id(PeerUser(id)),
+                utils.get_peer_id(PeerChat(id)),
+                utils.get_peer_id(PeerChannel(id))
+            )
+            return self._fetchone_entity(
+                'select id, hash from entities where id in (?,?,?)', ids
+            )
 
     # File processing