diff --git a/telethon/telegram_client.py b/telethon/telegram_client.py index 8546c377..32ade1a9 100644 --- a/telethon/telegram_client.py +++ b/telethon/telegram_client.py @@ -1,5 +1,7 @@ +import itertools import os import time +from collections import OrderedDict from datetime import datetime, timedelta from mimetypes import guess_type @@ -16,7 +18,7 @@ from .errors import ( ) from .network import ConnectionMode from .tl import TLObject -from .tl.custom import Draft +from .tl.custom import Draft, Dialog from .tl.entity_database import EntityDatabase from .tl.functions.account import ( GetPasswordRequest @@ -294,15 +296,14 @@ class TelegramClient(TelegramBareClient): The message ID to be used as an offset. :param offset_peer: The peer to be used as an offset. - :return: A tuple of lists ([dialogs], [entities]). + + :return List[telethon.tl.custom.Dialog]: A list dialogs. """ limit = float('inf') if limit is None else int(limit) if limit == 0: return [], [] - dialogs = {} # Use peer id as identifier to avoid dupes - messages = {} # Used later for sorting TODO also return these? - entities = {} + dialogs = OrderedDict() # Use peer id as identifier to avoid dupes while len(dialogs) < limit: real_limit = min(limit - len(dialogs), 100) r = self(GetDialogsRequest( @@ -312,16 +313,13 @@ class TelegramClient(TelegramBareClient): limit=real_limit )) - for d in r.dialogs: - dialogs[utils.get_peer_id(d.peer, True)] = d - for m in r.messages: - messages[m.id] = m + messages = {m.id: m for m in r.messages} + entities = {utils.get_peer_id(x, add_mark=True): x + for x in itertools.chain(r.users, r.chats)} - # We assume users can't have the same ID as a chat - for u in r.users: - entities[u.id] = u - for c in r.chats: - entities[c.id] = c + for d in r.dialogs: + dialogs[utils.get_peer_id(d.peer, add_mark=True)] = \ + Dialog(self, d, entities, messages) if len(r.dialogs) < real_limit or not isinstance(r, DialogsSlice): # Less than we requested means we reached the end, or @@ -334,20 +332,8 @@ class TelegramClient(TelegramBareClient): ) offset_id = r.messages[-1].id & 4294967296 # Telegram/danog magic - # Sort by message date. Windows will raise if timestamp is 0, - # so we need to set at least one day ahead while still being - # the smallest date possible. - no_date = datetime.fromtimestamp(86400) - ds = list(sorted( - dialogs.values(), - key=lambda d: getattr(messages[d.top_message], 'date', no_date) - )) - if limit < float('inf'): - ds = ds[:limit] - return ( - ds, - [utils.find_user_or_chat(d.peer, entities, entities) for d in ds] - ) + dialogs = list(dialogs.values()) + return dialogs[:limit] if limit < float('inf') else dialogs def get_drafts(self): # TODO: Ability to provide a `filter` """ diff --git a/telethon/tl/custom/__init__.py b/telethon/tl/custom/__init__.py index 40914f16..5b6bf44d 100644 --- a/telethon/tl/custom/__init__.py +++ b/telethon/tl/custom/__init__.py @@ -1 +1,2 @@ from .draft import Draft +from .dialog import Dialog diff --git a/telethon/tl/custom/dialog.py b/telethon/tl/custom/dialog.py new file mode 100644 index 00000000..bac8b0de --- /dev/null +++ b/telethon/tl/custom/dialog.py @@ -0,0 +1,37 @@ +from . import Draft +from ... import utils + + +class Dialog: + """ + Custom class that encapsulates a dialog (an open "conversation" with + someone, a group or a channel) providing an abstraction to easily + access the input version/normal entity/message etc. The library will + return instances of this class when calling `client.get_dialogs()`. + """ + def __init__(self, client, dialog, entities, messages): + # Both entities and messages being dicts {ID: item} + self._client = client + self.dialog = dialog + self.pinned = bool(dialog.pinned) + self.message = messages.get(dialog.top_message, None) + self.date = getattr(self.message, 'date', None) + + self.entity = entities[utils.get_peer_id(dialog.peer, add_mark=True)] + self.input_entity = utils.get_input_peer(self.entity) + self.name = utils.get_display_name(self.entity) + + self.unread_count = dialog.unread_count + self.unread_mentions_count = dialog.unread_mentions_count + + if dialog.draft: + self.draft = Draft(client, dialog.peer, dialog.draft) + else: + self.draft = None + + def send_message(self, *args, **kwargs): + """ + Sends a message to this dialog. This is just a wrapper around + client.send_message(dialog.input_entity, *args, **kwargs). + """ + return self._client.send_message(self.input_entity, *args, **kwargs) diff --git a/telethon/utils.py b/telethon/utils.py index 3259c8e2..5e92b13d 100644 --- a/telethon/utils.py +++ b/telethon/utils.py @@ -35,12 +35,12 @@ def get_display_name(entity): elif entity.last_name: return entity.last_name else: - return '(No name)' + return '' - if isinstance(entity, (Chat, Channel)): + elif isinstance(entity, (Chat, Channel)): return entity.title - return '(unknown)' + return '' # For some reason, .webp (stickers' format) is not registered add_type('image/webp', '.webp')