Fix entities weren't being passed to Draft

This commit is contained in:
Lonami Exo 2019-06-28 20:34:30 +02:00
parent 81e628b9f7
commit 84c4fcdec6
2 changed files with 8 additions and 3 deletions

View File

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

View File

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