Create a convenient class to wrap Dialog instances

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
+        dialogs = OrderedDict()  # Use peer id as identifier to avoid dupes
-        messages = {}  # Used later for sorting TODO also return these?
-        entities = {}
         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:
+            messages = {m.id: m for m in r.messages}
-                dialogs[utils.get_peer_id(d.peer, True)] = d
+            entities = {utils.get_peer_id(x, add_mark=True): x
-            for m in r.messages:
+                        for x in itertools.chain(r.users, r.chats)}
-                messages[m.id] = m
 
-            # We assume users can't have the same ID as a chat
+            for d in r.dialogs:
-            for u in r.users:
+                dialogs[utils.get_peer_id(d.peer, add_mark=True)] = \
-                entities[u.id] = u
+                    Dialog(self, d, entities, messages)
-            for c in r.chats:
-                entities[c.id] = c
 
             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,
+        dialogs = list(dialogs.values())
-        # so we need to set at least one day ahead while still being
+        return dialogs[:limit] if limit < float('inf') else dialogs
-        # 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]
-        )
 
     def get_drafts(self):  # TODO: Ability to provide a `filter`
         """

telethon/tl/custom/__init__.py

@@ -1 +1,2 @@
 from .draft import Draft
+from .dialog import Dialog

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)

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')