|
16 | 16 | # You should have received a copy of the GNU Lesser General Public License |
17 | 17 | # along with Pyrogram. If not, see <http://www.gnu.org/licenses/>. |
18 | 18 |
|
19 | | -import logging |
20 | | -from datetime import datetime |
21 | | -from typing import List |
| 19 | +from typing import AsyncGenerator, Optional |
22 | 20 |
|
23 | 21 | import pyrogram |
24 | | -from pyrogram import raw |
25 | | -from pyrogram import types |
26 | | -from pyrogram import utils |
27 | | - |
28 | | -log = logging.getLogger(__name__) |
| 22 | +from pyrogram import types, raw, utils |
29 | 23 |
|
30 | 24 |
|
31 | 25 | class GetDialogs: |
32 | 26 | async def get_dialogs( |
33 | 27 | self: "pyrogram.Client", |
34 | | - offset_date: datetime = datetime.fromtimestamp(0), |
35 | | - limit: int = 100, |
36 | | - pinned_only: bool = False |
37 | | - ) -> List["types.Dialog"]: |
38 | | - """Get a chunk of the user's dialogs. |
39 | | -
|
40 | | - You can get up to 100 dialogs at once. |
41 | | - For a more convenient way of getting a user's dialogs see :meth:`~pyrogram.Client.iter_dialogs`. |
| 28 | + limit: int = 0 |
| 29 | + ) -> Optional[AsyncGenerator["types.Dialog", None]]: |
| 30 | + """Get a user's dialogs sequentially. |
42 | 31 |
|
43 | 32 | Parameters: |
44 | | - offset_date (:py:obj:`~datetime.datetime`): |
45 | | - The offset date taken from the top message of a :obj:`~pyrogram.types.Dialog`. |
46 | | - Defaults to epoch. Valid for non-pinned dialogs only. |
47 | | -
|
48 | | - limit (``str``, *optional*): |
| 33 | + limit (``int``, *optional*): |
49 | 34 | Limits the number of dialogs to be retrieved. |
50 | | - Defaults to 100. Valid for non-pinned dialogs only. |
51 | | -
|
52 | | - pinned_only (``bool``, *optional*): |
53 | | - Pass True if you want to get only pinned dialogs. |
54 | | - Defaults to False. |
| 35 | + By default, no limit is applied and all dialogs are returned. |
55 | 36 |
|
56 | 37 | Returns: |
57 | | - List of :obj:`~pyrogram.types.Dialog`: On success, a list of dialogs is returned. |
| 38 | + ``Generator``: A generator yielding :obj:`~pyrogram.types.Dialog` objects. |
58 | 39 |
|
59 | 40 | Example: |
60 | 41 | .. code-block:: python |
61 | 42 |
|
62 | | - # Get first 100 dialogs |
63 | | - await app.get_dialogs() |
64 | | -
|
65 | | - # Get pinned dialogs |
66 | | - await app.get_dialogs(pinned_only=True) |
| 43 | + # Iterate through all dialogs |
| 44 | + async for dialog in app.get_dialogs(): |
| 45 | + print(dialog.chat.first_name or dialog.chat.title) |
67 | 46 | """ |
| 47 | + current = 0 |
| 48 | + total = limit or (1 << 31) - 1 |
| 49 | + limit = min(100, total) |
68 | 50 |
|
69 | | - if pinned_only: |
70 | | - r = await self.invoke( |
71 | | - raw.functions.messages.GetPinnedDialogs(folder_id=0), |
72 | | - sleep_threshold=60 |
73 | | - ) |
74 | | - else: |
| 51 | + offset_date = 0 |
| 52 | + offset_id = 0 |
| 53 | + offset_peer = raw.types.InputPeerEmpty() |
| 54 | + |
| 55 | + while True: |
75 | 56 | r = await self.invoke( |
76 | 57 | raw.functions.messages.GetDialogs( |
77 | | - offset_date=utils.datetime_to_timestamp(offset_date), |
78 | | - offset_id=0, |
79 | | - offset_peer=raw.types.InputPeerEmpty(), |
| 58 | + offset_date=offset_date, |
| 59 | + offset_id=offset_id, |
| 60 | + offset_peer=offset_peer, |
80 | 61 | limit=limit, |
81 | | - hash=0, |
82 | | - exclude_pinned=True |
| 62 | + hash=0 |
83 | 63 | ), |
84 | 64 | sleep_threshold=60 |
85 | 65 | ) |
86 | 66 |
|
87 | | - users = {i.id: i for i in r.users} |
88 | | - chats = {i.id: i for i in r.chats} |
| 67 | + users = {i.id: i for i in r.users} |
| 68 | + chats = {i.id: i for i in r.chats} |
| 69 | + |
| 70 | + messages = {} |
| 71 | + |
| 72 | + for message in r.messages: |
| 73 | + if isinstance(message, raw.types.MessageEmpty): |
| 74 | + continue |
| 75 | + |
| 76 | + chat_id = utils.get_peer_id(message.peer_id) |
| 77 | + messages[chat_id] = await types.Message._parse(self, message, users, chats) |
| 78 | + |
| 79 | + dialogs = [] |
| 80 | + |
| 81 | + for dialog in r.dialogs: |
| 82 | + if not isinstance(dialog, raw.types.Dialog): |
| 83 | + continue |
89 | 84 |
|
90 | | - messages = {} |
| 85 | + dialogs.append(types.Dialog._parse(self, dialog, messages, users, chats)) |
91 | 86 |
|
92 | | - for message in r.messages: |
93 | | - if isinstance(message, raw.types.MessageEmpty): |
94 | | - continue |
| 87 | + if not dialogs: |
| 88 | + return |
95 | 89 |
|
96 | | - chat_id = utils.get_peer_id(message.peer_id) |
97 | | - messages[chat_id] = await types.Message._parse(self, message, users, chats) |
| 90 | + last = dialogs[-1] |
98 | 91 |
|
99 | | - parsed_dialogs = [] |
| 92 | + offset_id = last.top_message.id |
| 93 | + offset_date = utils.datetime_to_timestamp(last.top_message.date) |
| 94 | + offset_peer = await self.resolve_peer(last.chat.id) |
100 | 95 |
|
101 | | - for dialog in r.dialogs: |
102 | | - if not isinstance(dialog, raw.types.Dialog): |
103 | | - continue |
| 96 | + for dialog in dialogs: |
| 97 | + yield dialog |
104 | 98 |
|
105 | | - parsed_dialogs.append(types.Dialog._parse(self, dialog, messages, users, chats)) |
| 99 | + current += 1 |
106 | 100 |
|
107 | | - return types.List(parsed_dialogs) |
| 101 | + if current >= total: |
| 102 | + return |
0 commit comments