LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2025-03-03 19:00:21 +03:00
Code Issues Packages Projects Releases Wiki Activity

Create client.iter_ versions for all client.get_ methods

Browse Source 
While doing so, the client.iter_drafts method has been simplified
as it made some unnecessary calls.

client.get_message_history has been shortened to client.get_messages,
and fixes a bug where the limit being zero made it return a tuple.

client.iter_messages also uses a local dictionary for entities so
it should become less big in memory (and possibly faster).

client.get_participants would fail with user entities, returning
only their input version.
This commit is contained in:
Lonami Exo 2018-03-08 11:44:13 +01:00
parent 09f0f86f1e
commit 5673866553
1 changed files with 158 additions and 101 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

245
telethon/telegram_client.py
View File

@ -90,6 +90,12 @@ from .extensions import markdown, html
__log__ = logging.getLogger(__name__) __log__ = logging.getLogger(__name__)
class _Box:
"""Helper class to pass parameters by reference"""
def __init__(self, x=None):
self.x = x
class TelegramClient(TelegramBareClient): class TelegramClient(TelegramBareClient):
""" """
Initializes the Telegram client with the specified API ID and Hash. Initializes the Telegram client with the specified API ID and Hash.
@ -508,10 +514,11 @@ class TelegramClient(TelegramBareClient):
# region Dialogs ("chats") requests # region Dialogs ("chats") requests
def get_dialogs(self, limit=10, offset_date=None, offset_id=0, def iter_dialogs(self, limit=None, offset_date=None, offset_id=0,
offset_peer=InputPeerEmpty()): offset_peer=InputPeerEmpty(), _total_box=None):
""" """
Gets N "dialogs" (open "chats" or conversations with other people). Returns an iterator over the dialogs, yielding 'limit' at most.
Dialogs are the open "chats" or conversations with other people.
Args: Args:
limit (:obj:`int` | :obj:`None`): limit (:obj:`int` | :obj:`None`):
@ -530,11 +537,16 @@ class TelegramClient(TelegramBareClient):
offset_peer (:obj:`InputPeer`, optional): offset_peer (:obj:`InputPeer`, optional):
The peer to be used as an offset. The peer to be used as an offset.
Returns: _total_box (:obj:`_Box`, optional):
A list dialogs, with an additional .total attribute on the list. A _Box instance to pass the total parameter by reference.
Yields:
Instances of ``telethon.tl.custom.Dialog``.
""" """
limit = float('inf') if limit is None else int(limit) limit = float('inf') if limit is None else int(limit)
if limit == 0: if limit == 0:
if not _total_box:
return
# Special case, get a single dialog and determine count # Special case, get a single dialog and determine count
dialogs = self(GetDialogsRequest( dialogs = self(GetDialogsRequest(
offset_date=offset_date, offset_date=offset_date,
@ -542,14 +554,12 @@ class TelegramClient(TelegramBareClient):
offset_peer=offset_peer, offset_peer=offset_peer,
limit=1 limit=1
)) ))
result = UserList() _total_box.x = getattr(dialogs, 'count', len(dialogs.dialogs))
result.total = getattr(dialogs, 'count', len(dialogs.dialogs)) return
return result
total_count = 0 seen = set()
dialogs = OrderedDict() # Use peer id as identifier to avoid dupes while len(seen) < limit:
while len(dialogs) < limit: real_limit = min(limit - len(seen), 100)
real_limit = min(limit - len(dialogs), 100)
r = self(GetDialogsRequest( r = self(GetDialogsRequest(
offset_date=offset_date, offset_date=offset_date,
offset_id=offset_id, offset_id=offset_id,
@ -557,14 +567,17 @@ class TelegramClient(TelegramBareClient):
limit=real_limit limit=real_limit
)) ))
total_count = getattr(r, 'count', len(r.dialogs)) if _total_box:
_total_box.x = getattr(r, 'count', len(r.dialogs))
messages = {m.id: m for m in r.messages} messages = {m.id: m for m in r.messages}
entities = {utils.get_peer_id(x): x entities = {utils.get_peer_id(x): x
for x in itertools.chain(r.users, r.chats)} for x in itertools.chain(r.users, r.chats)}
for d in r.dialogs: for d in r.dialogs:
dialogs[utils.get_peer_id(d.peer)] = \ peer_id = utils.get_peer_id(d.peer)
Dialog(self, d, entities, messages) if peer_id not in seen:
seen.add(peer_id)
yield Dialog(self, d, entities, messages)
if len(r.dialogs) < real_limit or not isinstance(r, DialogsSlice): if len(r.dialogs) < real_limit or not isinstance(r, DialogsSlice):
# Less than we requested means we reached the end, or # Less than we requested means we reached the end, or
@ -575,26 +588,33 @@ class TelegramClient(TelegramBareClient):
offset_peer = entities[utils.get_peer_id(r.dialogs[-1].peer)] offset_peer = entities[utils.get_peer_id(r.dialogs[-1].peer)]
offset_id = r.messages[-1].id offset_id = r.messages[-1].id
dialogs = UserList( def get_dialogs(self, *args, **kwargs):
itertools.islice(dialogs.values(), min(limit, len(dialogs))) """
) Same as :meth:`iter_dialogs`, but returns a list instead
dialogs.total = total_count with an additional .total attribute on the list.
"""
total_box = _Box(0)
kwargs['_total_box'] = total_box
dialogs = UserList(self.iter_dialogs(*args, **kwargs))
dialogs.total = total_box.x
return dialogs return dialogs
def get_drafts(self): # TODO: Ability to provide a `filter` def iter_drafts(self): # TODO: Ability to provide a `filter`
""" """
Gets all open draft messages. Iterator over all open draft messages.
Returns: The yielded items are custom ``Draft`` objects that are easier to use.
A list of custom ``Draft`` objects that are easy to work with:
You can call ``draft.set_message('text')`` to change the message, You can call ``draft.set_message('text')`` to change the message,
or delete it through :meth:`draft.delete()`. or delete it through :meth:`draft.delete()`.
""" """
response = self(GetAllDraftsRequest()) for update in self(GetAllDraftsRequest()).updates:
self.session.process_entities(response) yield Draft._from_update(self, update)
self.session.generate_sequence(response.seq)
drafts = [Draft._from_update(self, u) for u in response.updates] def get_drafts(self):
return drafts """
Same as :meth:`iter_drafts`, but returns a list instead.
"""
return list(self.iter_drafts())
@staticmethod @staticmethod
def _get_response_message(request, result): def _get_response_message(request, result):
@ -891,11 +911,11 @@ class TelegramClient(TelegramBareClient):
else: else:
return self(messages.DeleteMessagesRequest(message_ids, revoke=revoke)) return self(messages.DeleteMessagesRequest(message_ids, revoke=revoke))
def get_message_history(self, entity, limit=20, offset_date=None, def iter_messages(self, entity, limit=20, offset_date=None,
offset_id=0, max_id=0, min_id=0, add_offset=0, offset_id=0, max_id=0, min_id=0, add_offset=0,
batch_size=100, wait_time=None): batch_size=100, wait_time=None, _total_box=None):
""" """
Gets the message history for the specified entity Iterator over the message history for the specified entity.
Args: Args:
entity (:obj:`entity`): entity (:obj:`entity`):
@ -939,10 +959,12 @@ class TelegramClient(TelegramBareClient):
If left to ``None``, it will default to 1 second only if If left to ``None``, it will default to 1 second only if
the limit is higher than 3000. the limit is higher than 3000.
Returns: _total_box (:obj:`_Box`, optional):
A list of messages with extra attributes: A _Box instance to pass the total parameter by reference.
Yields:
Instances of ``telethon.tl.types.Message`` with extra attributes:
* ``.total`` = (on the list) total amount of messages sent.
* ``.sender`` = entity of the sender. * ``.sender`` = entity of the sender.
* ``.fwd_from.sender`` = if fwd_from, who sent it originally. * ``.fwd_from.sender`` = if fwd_from, who sent it originally.
* ``.fwd_from.channel`` = if fwd_from, original channel. * ``.fwd_from.channel`` = if fwd_from, original channel.
@ -959,25 +981,26 @@ class TelegramClient(TelegramBareClient):
entity = self.get_input_entity(entity) entity = self.get_input_entity(entity)
limit = float('inf') if limit is None else int(limit) limit = float('inf') if limit is None else int(limit)
if limit == 0: if limit == 0:
if not _total_box:
return
# No messages, but we still need to know the total message count # No messages, but we still need to know the total message count
result = self(GetHistoryRequest( result = self(GetHistoryRequest(
peer=entity, limit=1, peer=entity, limit=1,
offset_date=None, offset_id=0, max_id=0, min_id=0, offset_date=None, offset_id=0, max_id=0, min_id=0,
add_offset=0, hash=0 add_offset=0, hash=0
)) ))
return getattr(result, 'count', len(result.messages)), [], [] _total_box.x = getattr(result, 'count', len(result.messages))
return
if wait_time is None: if wait_time is None:
wait_time = 1 if limit > 3000 else 0 wait_time = 1 if limit > 3000 else 0
have = 0
batch_size = min(max(batch_size, 1), 100) batch_size = min(max(batch_size, 1), 100)
total_messages = 0 while have < limit:
messages = UserList()
entities = {}
while len(messages) < limit:
# Telegram has a hard limit of 100 # Telegram has a hard limit of 100
real_limit = min(limit - len(messages), batch_size) real_limit = min(limit - have, batch_size)
result = self(GetHistoryRequest( r = self(GetHistoryRequest(
peer=entity, peer=entity,
limit=real_limit, limit=real_limit,
offset_date=offset_date, offset_date=offset_date,
@ -987,48 +1010,63 @@ class TelegramClient(TelegramBareClient):
add_offset=add_offset, add_offset=add_offset,
hash=0 hash=0
)) ))
messages.extend( if _total_box:
m for m in result.messages if not isinstance(m, MessageEmpty) _total_box.x = getattr(r, 'count', len(r.messages))
)
total_messages = getattr(result, 'count', len(result.messages))
for u in result.users: entities = {utils.get_peer_id(x): x
entities[utils.get_peer_id(u)] = u for x in itertools.chain(r.users, r.chats)}
for c in result.chats:
entities[utils.get_peer_id(c)] = c
if len(result.messages) < real_limit: for message in r.messages:
break if isinstance(message, MessageEmpty):
continue
offset_id = result.messages[-1].id # Add a few extra attributes to the Message to be friendlier.
offset_date = result.messages[-1].date
time.sleep(wait_time)
# Add a few extra attributes to the Message to make it friendlier.
messages.total = total_messages
for m in messages:
# To make messages more friendly, always add message # To make messages more friendly, always add message
# to service messages, and action to normal messages. # to service messages, and action to normal messages.
m.message = getattr(m, 'message', None) message.message = getattr(message, 'message', None)
m.action = getattr(m, 'action', None) message.action = getattr(message, 'action', None)
m.sender = (None if not m.from_id else message.to = entities[utils.get_peer_id(message.to_id)]
entities[utils.get_peer_id(m.from_id)]) message.sender = (
None if not message.from_id else
if getattr(m, 'fwd_from', None): entities[utils.get_peer_id(message.from_id)]
m.fwd_from.sender = (
None if not m.fwd_from.from_id else
entities[utils.get_peer_id(m.fwd_from.from_id)]
) )
m.fwd_from.channel = ( if getattr(message, 'fwd_from', None):
None if not m.fwd_from.channel_id else message.fwd_from.sender = (
None if not message.fwd_from.from_id else
entities[utils.get_peer_id(message.fwd_from.from_id)]
)
message.fwd_from.channel = (
None if not message.fwd_from.channel_id else
entities[utils.get_peer_id( entities[utils.get_peer_id(
PeerChannel(m.fwd_from.channel_id) PeerChannel(message.fwd_from.channel_id)
)] )]
) )
yield message
have += 1
m.to = entities[utils.get_peer_id(m.to_id)] if len(r.messages) < real_limit:
break
return messages offset_id = r.messages[-1].id
offset_date = r.messages[-1].date
time.sleep(wait_time)
def get_messages(self, *args, **kwargs):
"""
Same as :meth:`iter_messages`, but returns a list instead
with an additional .total attribute on the list.
"""
total_box = _Box(0)
kwargs['_total_box'] = total_box
msgs = UserList(self.iter_messages(*args, **kwargs))
msgs.total = total_box.x
return msgs
def get_message_history(self, *args, **kwargs):
warnings.warn(
'get_message_history is deprecated, use get_messages instead'
)
return self.get_messages(*args, **kwargs)
def send_read_acknowledge(self, entity, message=None, max_id=None, def send_read_acknowledge(self, entity, message=None, max_id=None,
clear_mentions=False): clear_mentions=False):
@ -1096,8 +1134,8 @@ class TelegramClient(TelegramBareClient):
raise TypeError('Invalid message type: {}'.format(type(message))) raise TypeError('Invalid message type: {}'.format(type(message)))
def get_participants(self, entity, limit=None, search='', def iter_participants(self, entity, limit=None, search='',
aggressive=False): aggressive=False, _total_box=None):
""" """
Gets the list of participants from the specified entity. Gets the list of participants from the specified entity.
@ -1121,6 +1159,9 @@ class TelegramClient(TelegramBareClient):
This has no effect for groups or channels with less than This has no effect for groups or channels with less than
10,000 members. 10,000 members.
_total_box (:obj:`_Box`, optional):
A _Box instance to pass the total parameter by reference.
Returns: Returns:
A list of participants with an additional .total variable on the A list of participants with an additional .total variable on the
list indicating the total amount of members in this group/channel. list indicating the total amount of members in this group/channel.
@ -1131,12 +1172,13 @@ class TelegramClient(TelegramBareClient):
total = self(GetFullChannelRequest( total = self(GetFullChannelRequest(
entity entity
)).full_chat.participants_count )).full_chat.participants_count
if limit == 0: if _total_box:
users = UserList() _total_box.x = total
users.total = total
return users
all_participants = {} if limit == 0:
return
seen = set()
if total > 10000 and aggressive: if total > 10000 and aggressive:
requests = [GetParticipantsRequest( requests = [GetParticipantsRequest(
channel=entity, channel=entity,
@ -1176,25 +1218,40 @@ class TelegramClient(TelegramBareClient):
else: else:
requests[i].offset += len(participants.users) requests[i].offset += len(participants.users)
for user in participants.users: for user in participants.users:
if len(all_participants) < limit: if user.id not in seen:
all_participants[user.id] = user seen.add(user.id)
if limit < float('inf'): yield user
values = itertools.islice(all_participants.values(), limit) if len(seen) >= limit:
else: return
values = all_participants.values()
users = UserList(values)
users.total = total
elif isinstance(entity, InputPeerChat): elif isinstance(entity, InputPeerChat):
users = self(GetFullChatRequest(entity.chat_id)).users users = self(GetFullChatRequest(entity.chat_id)).users
if len(users) > limit: if _total_box:
users = users[:limit] _total_box.x = len(users)
users = UserList(users)
users.total = len(users) have = 0
for user in users:
have += 1
if have > limit:
break
else: else:
users = UserList(None if limit == 0 else [entity]) yield user
users.total = 1 else:
return users if _total_box:
_total_box.x = 1
if limit != 0:
yield self.get_entity(entity)
def get_participants(self, *args, **kwargs):
"""
Same as :meth:`iter_participants`, but returns a list instead
with an additional .total attribute on the list.
"""
total_box = _Box(0)
kwargs['_total_box'] = total_box
dialogs = UserList(self.iter_participants(*args, **kwargs))
dialogs.total = total_box.x
return dialogs
# endregion # endregion