Add support for "send_as" chats

delivrance · delivrance · commit 9a2bc25bc7f3 · 2021-12-23T16:53:03.000+01:00
- Add methods get_send_as_chats() and set_send_as_chat()
- Add field Chat.send_as_chat
diff --git a/compiler/docs/compiler.py b/compiler/docs/compiler.py
@@ -224,6 +224,8 @@ def get_title_list(s: str) -> list:
             mark_chat_unread
             get_chat_event_log
             get_chat_online_count
+            get_send_as_chats
+            set_send_as_chat
         """,
         users="""
         Users
diff --git a/pyrogram/methods/chats/__init__.py b/pyrogram/methods/chats/__init__.py
@@ -35,6 +35,7 @@
 from .get_dialogs import GetDialogs
 from .get_dialogs_count import GetDialogsCount
 from .get_nearby_chats import GetNearbyChats
+from .get_send_as_chats import GetSendAsChats
 from .iter_chat_members import IterChatMembers
 from .iter_dialogs import IterDialogs
 from .join_chat import JoinChat
@@ -48,6 +49,7 @@
 from .set_chat_permissions import SetChatPermissions
 from .set_chat_photo import SetChatPhoto
 from .set_chat_title import SetChatTitle
+from .set_send_as_chat import SetSendAsChat
 from .set_slow_mode import SetSlowMode
 from .unarchive_chats import UnarchiveChats
 from .unban_chat_member import UnbanChatMember
@@ -94,6 +96,8 @@ class Chats(
     UnpinAllChatMessages,
     MarkChatUnread,
     GetChatEventLog,
-    GetChatOnlineCount
+    GetChatOnlineCount,
+    GetSendAsChats,
+    SetSendAsChat
 ):
     pass
diff --git a/pyrogram/methods/chats/get_send_as_chats.py b/pyrogram/methods/chats/get_send_as_chats.py
@@ -0,0 +1,63 @@
+#  Pyrogram - Telegram MTProto API Client Library for Python
+#  Copyright (C) 2017-2021 Dan <https://github.com/delivrance>
+#
+#  This file is part of Pyrogram.
+#
+#  Pyrogram is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU Lesser General Public License as published
+#  by the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  Pyrogram is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU Lesser General Public License for more details.
+#
+#  You should have received a copy of the GNU Lesser General Public License
+#  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+from typing import List, Union
+
+from pyrogram import raw
+from pyrogram import types
+from pyrogram.scaffold import Scaffold
+
+
+class GetSendAsChats(Scaffold):
+    async def get_send_as_chats(
+        self,
+        chat_id: Union[int, str]
+    ) -> List["types.Chat"]:
+        """Get the list of "send_as" chats available.
+
+        Parameters:
+            chat_id (``int`` | ``str``):
+                Unique identifier (int) or username (str) of the target chat.
+
+        Returns:
+            List[:obj:`~pyrogram.types.Chat`]: The list of chats.
+
+        Example:
+            .. code-block:: python
+
+                chats = app.get_send_as_chats(chat_id)
+                print(chats)
+        """
+        r = await self.send(
+            raw.functions.channels.GetSendAs(
+                peer=await self.resolve_peer(chat_id)
+            )
+        )
+
+        users = {u.id: u for u in r.users}
+        chats = {c.id: c for c in r.chats}
+
+        send_as_chats = types.List()
+
+        for p in r.peers:
+            if isinstance(p, raw.types.PeerUser):
+                send_as_chats.append(types.Chat._parse_chat(self, users[p.user_id]))
+            else:
+                send_as_chats.append(types.Chat._parse_chat(self, chats[p.channel_id]))
+
+        return send_as_chats
diff --git a/pyrogram/methods/chats/set_send_as_chat.py b/pyrogram/methods/chats/set_send_as_chat.py
@@ -0,0 +1,55 @@
+#  Pyrogram - Telegram MTProto API Client Library for Python
+#  Copyright (C) 2017-2021 Dan <https://github.com/delivrance>
+#
+#  This file is part of Pyrogram.
+#
+#  Pyrogram is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU Lesser General Public License as published
+#  by the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  Pyrogram is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU Lesser General Public License for more details.
+#
+#  You should have received a copy of the GNU Lesser General Public License
+#  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
+
+from typing import Union
+
+from pyrogram import raw
+from pyrogram.scaffold import Scaffold
+
+
+class SetSendAsChat(Scaffold):
+    async def set_send_as_chat(
+        self,
+        chat_id: Union[int, str],
+        send_as_chat_id: Union[int, str]
+    ) -> bool:
+        """Set the default "send_as" chat for a chat.
+
+        Use :meth:`~pyrogram.Client.get_send_as_chats` to get all the "send_as" chats available for use.
+
+        Parameters:
+            chat_id (``int`` | ``str``):
+                Unique identifier (int) or username (str) of the target chat.
+
+            send_as_chat_id (``int`` | ``str``):
+                Unique identifier (int) or username (str) of the send_as chat.
+
+        Returns:
+            ``bool``: On success, true is returned
+
+        Example:
+            .. code-block:: python
+
+                app.set_send_as_chat(chat_id, send_as_chat_id)
+        """
+        return await self.send(
+            raw.functions.messages.SaveDefaultSendAs(
+                peer=await self.resolve_peer(chat_id),
+                send_as=await self.resolve_peer(send_as_chat_id)
+            )
+        )
diff --git a/pyrogram/types/user_and_chats/chat.py b/pyrogram/types/user_and_chats/chat.py
@@ -120,6 +120,10 @@ class Chat(Object):
         linked_chat (:obj:`~pyrogram.types.Chat`, *optional*):
             The linked discussion group (in case of channels) or the linked channel (in case of supergroups).
             Returned only in :meth:`~pyrogram.Client.get_chat`.
+
+        send_as_chat (:obj:`~pyrogram.types.Chat`, *optional*):
+            The default "send_as" chat.
+            Returned only in :meth:`~pyrogram.Client.get_chat`.
     """
 
     def __init__(
@@ -151,7 +155,8 @@ def __init__(
         restrictions: List["types.Restriction"] = None,
         permissions: "types.ChatPermissions" = None,
         distance: int = None,
-        linked_chat: "types.Chat" = None
+        linked_chat: "types.Chat" = None,
+        send_as_chat: "types.Chat" = None
     ):
         super().__init__(client)
 
@@ -181,6 +186,7 @@ def __init__(
         self.permissions = permissions
         self.distance = distance
         self.linked_chat = linked_chat
+        self.send_as_chat = send_as_chat
 
     @staticmethod
     def _parse_user_chat(client, user: raw.types.User) -> "Chat":
@@ -275,43 +281,51 @@ def _parse_dialog(client, peer, users: dict, chats: dict):
 
     @staticmethod
     async def _parse_full(client, chat_full: Union[raw.types.messages.ChatFull, raw.types.users.UserFull]) -> "Chat":
+        users = {u.id: u for u in chat_full.users}
+        chats = {c.id: c for c in chat_full.chats}
+
         if isinstance(chat_full, raw.types.users.UserFull):
-            parsed_chat = Chat._parse_user_chat(client, chat_full.users[0])
-            parsed_chat.bio = chat_full.full_user.about
+            full_user = chat_full.full_user
 
-            if chat_full.full_user.pinned_msg_id:
+            parsed_chat = Chat._parse_user_chat(client, users[full_user.id])
+            parsed_chat.bio = full_user.about
+
+            if full_user.pinned_msg_id:
                 parsed_chat.pinned_message = await client.get_messages(
                     parsed_chat.id,
-                    message_ids=chat_full.full_user.pinned_msg_id
+                    message_ids=full_user.pinned_msg_id
                 )
         else:
             full_chat = chat_full.full_chat
-            chat = None
-            linked_chat = None
-
-            for c in chat_full.chats:
-                if full_chat.id == c.id:
-                    chat = c
-
-                if isinstance(full_chat, raw.types.ChannelFull):
-                    if full_chat.linked_chat_id == c.id:
-                        linked_chat = c
+            chat_raw = chats[full_chat.id]
 
             if isinstance(full_chat, raw.types.ChatFull):
-                parsed_chat = Chat._parse_chat_chat(client, chat)
+                parsed_chat = Chat._parse_chat_chat(client, chat_raw)
                 parsed_chat.description = full_chat.about or None
 
                 if isinstance(full_chat.participants, raw.types.ChatParticipants):
                     parsed_chat.members_count = len(full_chat.participants.participants)
             else:
-                parsed_chat = Chat._parse_channel_chat(client, chat)
+                parsed_chat = Chat._parse_channel_chat(client, chat_raw)
                 parsed_chat.members_count = full_chat.participants_count
                 parsed_chat.description = full_chat.about or None
                 # TODO: Add StickerSet type
                 parsed_chat.can_set_sticker_set = full_chat.can_set_stickers
                 parsed_chat.sticker_set_name = getattr(full_chat.stickerset, "short_name", None)
-                if linked_chat:
-                    parsed_chat.linked_chat = Chat._parse_channel_chat(client, linked_chat)
+
+                linked_chat_raw = chats.get(full_chat.linked_chat_id, None)
+
+                if linked_chat_raw:
+                    parsed_chat.linked_chat = Chat._parse_channel_chat(client, linked_chat_raw)
+
+                default_send_as = full_chat.default_send_as
+
+                if isinstance(default_send_as, raw.types.PeerUser):
+                    send_as_raw = users[default_send_as.user_id]
+                else:
+                    send_as_raw = chats[default_send_as.channel_id]
+
+                parsed_chat.send_as_chat = Chat._parse_chat(client, send_as_raw)
 
             if full_chat.pinned_msg_id:
                 parsed_chat.pinned_message = await client.get_messages(
