LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2024-11-26 11:23:46 +03:00
Code Issues Packages Projects Releases Wiki Activity

Support getting any entity by just their positive ID

Browse Source
This commit is contained in:
Lonami Exo 2018-03-08 10:05:40 +01:00
parent d3d190f36e
commit ce0dee63b1
4 changed files with 74 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
telethon/sessions/memory.py
View File

@ -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:
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)

23
telethon/sessions/sqlalchemy.py
View File

@ -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):

15
telethon/sessions/sqlite.py
View File

@ -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):
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

24
telethon/telegram_client.py
View File

@ -2272,23 +2272,21 @@ class TelegramClient(TelegramBareClient):
return InputPeerSelf()
return utils.get_input_peer(self._get_entity_from_string(peer))
if isinstance(peer, int):
peer, kind = utils.resolve_id(peer)
peer = kind(peer)
if not isinstance(peer, int):
try:
is_peer = peer.SUBCLASS_OF_ID == 0x2d45687 # crc32(b'Peer')
if not is_peer:
if peer.SUBCLASS_OF_ID != 0x2d45687: # crc32(b'Peer')
return utils.get_input_peer(peer)
except (AttributeError, TypeError):
is_peer = False
peer = None
if not is_peer:
if not peer:
raise TypeError(
'Cannot turn "{}" into an input entity.'.format(peer)
)
# Not found, look in the dialogs with the hope to find it.
# Add the mark to the peers if the user passed a Peer (not an int)
# Look in the dialogs with the hope to find it.
mark = not isinstance(peer, int)
target_id = utils.get_peer_id(peer)
req = GetDialogsRequest(
offset_date=None,
@ -2299,12 +2297,20 @@ class TelegramClient(TelegramBareClient):
while True:
result = self(req)
entities = {}
if mark:
for x in itertools.chain(result.users, result.chats):
x_id = utils.get_peer_id(x)
if x_id == target_id:
return utils.get_input_peer(x)
else:
entities[x_id] = x
else:
for x in itertools.chain(result.users, result.chats):
if x.id == target_id:
return utils.get_input_peer(x)
else:
entities[utils.get_peer_id(x)] = x
if len(result.dialogs) < req.limit:
break