From 84c4fcdec69a44bfd31149afd1c7366d6aedbd67 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Fri, 28 Jun 2019 20:34:30 +0200
Subject: [PATCH] Fix entities weren't being passed to Draft

---
 telethon/client/dialogs.py  | 7 ++++++-
 telethon/tl/custom/draft.py | 4 ++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/telethon/client/dialogs.py b/telethon/client/dialogs.py
index c388ac52..2fa2780b 100644
--- a/telethon/client/dialogs.py
+++ b/telethon/client/dialogs.py
@@ -96,7 +96,12 @@ class _DialogsIter(RequestIter):
 class _DraftsIter(RequestIter):
     async def _init(self, **kwargs):
         r = await self.client(functions.messages.GetAllDraftsRequest())
-        self.buffer.extend(custom.Draft._from_update(self.client, u)
+
+        # TODO Maybe there should be a helper method for this?
+        entities = {utils.get_peer_id(x): x
+                    for x in itertools.chain(r.users, r.chats)}
+
+        self.buffer.extend(custom.Draft._from_update(self.client, u, entities)
                            for u in r.updates)
 
     async def _load_next_chunk(self):
diff --git a/telethon/tl/custom/draft.py b/telethon/tl/custom/draft.py
index eb12ad3d..c5d5cc89 100644
--- a/telethon/tl/custom/draft.py
+++ b/telethon/tl/custom/draft.py
@@ -45,10 +45,10 @@ class Draft:
                    draft=dialog.dialog.draft, entity=dialog.entity)
 
     @classmethod
-    def _from_update(cls, client, update, entities=None):
+    def _from_update(cls, client, update, entities):
         assert isinstance(update, UpdateDraftMessage)
         return cls(client=client, peer=update.peer, draft=update.draft,
-                   entity=(entities or {}).get(get_peer_id(update.peer)))
+                   entity=entities.get(get_peer_id(update.peer)))
 
     @property
     def entity(self):