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