Add search/filter/from_user parameters to iter_messages

This commit is contained in:
Lonami Exo 2018-04-22 16:30:14 +02:00
parent 544651caa7
commit 6ecef42ec7

View File

@ -58,7 +58,7 @@ from .tl.functions.messages import (
SendMessageRequest, GetChatsRequest, GetAllDraftsRequest, SendMessageRequest, GetChatsRequest, GetAllDraftsRequest,
CheckChatInviteRequest, ReadMentionsRequest, SendMultiMediaRequest, CheckChatInviteRequest, ReadMentionsRequest, SendMultiMediaRequest,
UploadMediaRequest, EditMessageRequest, GetFullChatRequest, UploadMediaRequest, EditMessageRequest, GetFullChatRequest,
ForwardMessagesRequest ForwardMessagesRequest, SearchRequest
) )
from .tl.functions import channels from .tl.functions import channels
@ -990,10 +990,14 @@ class TelegramClient(TelegramBareClient):
def iter_messages(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,
search=None, filter=None, from_user=None,
batch_size=100, wait_time=None, _total=None): batch_size=100, wait_time=None, _total=None):
""" """
Iterator over the message history for the specified entity. Iterator over the message history for the specified entity.
If either `search`, `filter` or `from_user` are provided,
:tl:`messages.Search` will be used instead of :tl:`messages.getHistory`.
Args: Args:
entity (`entity`): entity (`entity`):
The entity from whom to retrieve the message history. The entity from whom to retrieve the message history.
@ -1025,6 +1029,17 @@ class TelegramClient(TelegramBareClient):
Additional message offset (all of the specified offsets + Additional message offset (all of the specified offsets +
this offset = older messages). this offset = older messages).
search (`str`):
The string to be used as a search query.
filter (:tl:`MessagesFilter` | `type`):
The filter to use when returning messages. For instance,
:tl:`InputMessagesFilterPhotos` would yield only messages
containing photos.
from_user (`entity`):
Only messages from this user will be returned.
batch_size (`int`): batch_size (`int`):
Messages will be returned in chunks of this size (100 is Messages will be returned in chunks of this size (100 is
the maximum). While it makes no sense to modify this value, the maximum). While it makes no sense to modify this value,
@ -1056,15 +1071,37 @@ 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 search is not None or filter or from_user:
request = SearchRequest(
peer=entity,
q=search or '',
filter=filter() if isinstance(filter, type) else filter,
min_date=None,
max_date=offset_date,
offset_id=offset_id,
add_offset=add_offset,
limit=1,
max_id=max_id,
min_id=min_id,
from_id=self.get_input_entity(from_user) if from_user else None
)
else:
request = GetHistoryRequest(
peer=entity,
limit=1,
offset_date=offset_date,
offset_id=offset_id,
min_id=min_id,
max_id=max_id,
add_offset=add_offset,
hash=0
)
if limit == 0: if limit == 0:
if not _total: if not _total:
return 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(request)
peer=entity, limit=1,
offset_date=None, offset_id=0, max_id=0, min_id=0,
add_offset=0, hash=0
))
_total[0] = getattr(result, 'count', len(result.messages)) _total[0] = getattr(result, 'count', len(result.messages))
return return
@ -1075,17 +1112,8 @@ class TelegramClient(TelegramBareClient):
batch_size = min(max(batch_size, 1), 100) batch_size = min(max(batch_size, 1), 100)
while have < limit: while have < limit:
# Telegram has a hard limit of 100 # Telegram has a hard limit of 100
real_limit = min(limit - have, batch_size) request.limit = min(limit - have, batch_size)
r = self(GetHistoryRequest( r = self(request)
peer=entity,
limit=real_limit,
offset_date=offset_date,
offset_id=offset_id,
max_id=max_id,
min_id=min_id,
add_offset=add_offset,
hash=0
))
if _total: if _total:
_total[0] = getattr(r, 'count', len(r.messages)) _total[0] = getattr(r, 'count', len(r.messages))
@ -1120,11 +1148,15 @@ class TelegramClient(TelegramBareClient):
yield message yield message
have += 1 have += 1
if len(r.messages) < real_limit: if len(r.messages) < request.limit:
break break
offset_id = r.messages[-1].id request.offset_id = r.messages[-1].id
offset_date = r.messages[-1].date if isinstance(request, GetHistoryRequest):
request.offset_date = r.messages[-1].date
else:
request.max_date = r.messages[-1].date
time.sleep(wait_time) time.sleep(wait_time)
def get_messages(self, *args, **kwargs): def get_messages(self, *args, **kwargs):