Return a UserList with a .total attribute for get dialogs/history

This commit is contained in:
Lonami Exo 2017-12-28 11:49:35 +01:00
parent 0755bda220
commit 459022bdab

View File

@ -1,7 +1,7 @@
import itertools import itertools
import os import os
import time import time
from collections import OrderedDict from collections import OrderedDict, UserList
from datetime import datetime, timedelta from datetime import datetime, timedelta
from mimetypes import guess_type from mimetypes import guess_type
@ -296,12 +296,23 @@ class TelegramClient(TelegramBareClient):
:param offset_peer: :param offset_peer:
The peer to be used as an offset. The peer to be used as an offset.
:return List[telethon.tl.custom.Dialog]: A list dialogs. :return UserList[telethon.tl.custom.Dialog]:
A list dialogs, with an additional .total attribute on the list.
""" """
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:
return [] # Special case, get a single dialog and determine count
dialogs = self(GetDialogsRequest(
offset_date=offset_date,
offset_id=offset_id,
offset_peer=offset_peer,
limit=1
))
result = UserList()
result.total = getattr(dialogs, 'count', len(dialogs.dialogs))
return result
total_count = 0
dialogs = OrderedDict() # Use peer id as identifier to avoid dupes dialogs = OrderedDict() # Use peer id as identifier to avoid dupes
while len(dialogs) < limit: while len(dialogs) < limit:
real_limit = min(limit - len(dialogs), 100) real_limit = min(limit - len(dialogs), 100)
@ -312,6 +323,7 @@ class TelegramClient(TelegramBareClient):
limit=real_limit limit=real_limit
)) ))
total_count = 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, add_mark=True): x entities = {utils.get_peer_id(x, add_mark=True): x
for x in itertools.chain(r.users, r.chats)} for x in itertools.chain(r.users, r.chats)}
@ -331,7 +343,8 @@ class TelegramClient(TelegramBareClient):
) )
offset_id = r.messages[-1].id & 4294967296 # Telegram/danog magic offset_id = r.messages[-1].id & 4294967296 # Telegram/danog magic
dialogs = list(dialogs.values()) dialogs = UserList(dialogs.values())
dialogs.total = total_count
return dialogs[:limit] if limit < float('inf') else dialogs return dialogs[:limit] if limit < float('inf') else dialogs
def get_drafts(self): # TODO: Ability to provide a `filter` def get_drafts(self): # TODO: Ability to provide a `filter`
@ -485,7 +498,7 @@ class TelegramClient(TelegramBareClient):
(all of the specified offsets + this offset = older messages). (all of the specified offsets + this offset = older messages).
:return: A list of messages with extra attributes: :return: A list of messages with extra attributes:
.total = total amount of messages in this history .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
@ -502,7 +515,7 @@ class TelegramClient(TelegramBareClient):
return getattr(result, 'count', len(result.messages)), [], [] return getattr(result, 'count', len(result.messages)), [], []
total_messages = 0 total_messages = 0
messages = [] messages = UserList()
entities = {} entities = {}
while len(messages) < limit: while len(messages) < limit:
# Telegram has a hard limit of 100 # Telegram has a hard limit of 100
@ -542,9 +555,9 @@ class TelegramClient(TelegramBareClient):
time.sleep(1) time.sleep(1)
# Add a few extra attributes to the Message to make it friendlier. # Add a few extra attributes to the Message to make it friendlier.
messages.total = total_messages
for m in messages: for m in messages:
# TODO Better way to return a total without tuples? # TODO Better way to return a total without tuples?
m.total = total_messages
m.sender = (None if not m.from_id else m.sender = (None if not m.from_id else
entities[utils.get_peer_id(m.from_id, add_mark=True)]) entities[utils.get_peer_id(m.from_id, add_mark=True)])