Update iter_dialogs.py

Closes #749 #750 #756
This commit is contained in:
Dan 2021-09-15 17:55:01 +02:00 committed by GitHub
parent 02a3969101
commit 09c8289259
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -18,15 +18,14 @@
from typing import AsyncGenerator, Optional from typing import AsyncGenerator, Optional
from pyrogram import types from pyrogram import types, raw, utils
from pyrogram.scaffold import Scaffold from pyrogram.scaffold import Scaffold
class IterDialogs(Scaffold): class IterDialogs(Scaffold):
async def iter_dialogs( async def iter_dialogs(
self, self,
limit: int = 0, limit: int = 0
offset_date: int = 0
) -> Optional[AsyncGenerator["types.Dialog", None]]: ) -> Optional[AsyncGenerator["types.Dialog", None]]:
"""Iterate through a user's dialogs sequentially. """Iterate through a user's dialogs sequentially.
@ -39,10 +38,6 @@ class IterDialogs(Scaffold):
Limits the number of dialogs to be retrieved. Limits the number of dialogs to be retrieved.
By default, no limit is applied and all dialogs are returned. By default, no limit is applied and all dialogs are returned.
offset_date (``int``):
The offset date in Unix time taken from the top message of a :obj:`~pyrogram.types.Dialog`.
Defaults to 0 (most recent dialog).
Returns: Returns:
``Generator``: A generator yielding :obj:`~pyrogram.types.Dialog` objects. ``Generator``: A generator yielding :obj:`~pyrogram.types.Dialog` objects.
@ -57,28 +52,50 @@ class IterDialogs(Scaffold):
total = limit or (1 << 31) - 1 total = limit or (1 << 31) - 1
limit = min(100, total) limit = min(100, total)
pinned_dialogs = await self.get_dialogs( offset_date = 0
pinned_only=True offset_id = 0
) offset_peer = raw.types.InputPeerEmpty()
for dialog in pinned_dialogs:
yield dialog
current += 1
if current >= total:
return
while True: while True:
dialogs = await self.get_dialogs( r = (await self.send(
offset_date=offset_date, raw.functions.messages.GetDialogs(
limit=limit offset_date=offset_date,
) offset_id=offset_id,
offset_peer=offset_peer,
limit=limit,
hash=0
),
sleep_threshold=60
))
users = {i.id: i for i in r.users}
chats = {i.id: i for i in r.chats}
messages = {}
for message in r.messages:
if isinstance(message, raw.types.MessageEmpty):
continue
chat_id = utils.get_peer_id(message.peer_id)
messages[chat_id] = await types.Message._parse(self, message, users, chats)
dialogs = []
for dialog in r.dialogs:
if not isinstance(dialog, raw.types.Dialog):
continue
dialogs.append(types.Dialog._parse(self, dialog, messages, users, chats))
if not dialogs: if not dialogs:
return return
offset_date = dialogs[-1].top_message.date last = dialogs[-1]
offset_id = last.top_message.message_id
offset_date = last.top_message.date
offset_peer = await self.resolve_peer(last.chat.id)
for dialog in dialogs: for dialog in dialogs:
yield dialog yield dialog