fix ChatMember parser

z44d · Jul 11, 2024 · 2553322 · 2553322
1 parent 1f02e9f
commit 2553322
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 26 deletions.
diff --git a/tgram/methods.py b/tgram/methods.py
@@ -26,6 +26,12 @@
     ChatInviteLink,
     ChatFullInfo,
     ChatMember,
+    ChatMemberOwner,
+    ChatMemberAdministrator,
+    ChatMemberBanned,
+    ChatMemberLeft,
+    ChatMemberRestricted,
+    ChatMemberMember,
     Sticker,
     ForumTopic,
     UserChatBoosts,
@@ -1263,7 +1269,7 @@ async def get_chat(self: "tgram.TgBot", chat_id: Union[int, str]) -> ChatFullInf
 
     async def get_chat_administrators(
         self: "tgram.TgBot", chat_id: Union[int, str]
-    ) -> List[ChatMember]:
+    ) -> List[Union[ChatMemberAdministrator, ChatMemberOwner]]:
         """https://core.telegram.org/bots/api/#getchatadministrators"""
         result = await self._send_request(
             "getChatAdministrators",
@@ -1283,7 +1289,14 @@ async def get_chat_member_count(
 
     async def get_chat_member(
         self: "tgram.TgBot", chat_id: Union[int, str], user_id: int
-    ) -> ChatMember:
+    ) -> Union[
+        ChatMemberOwner,
+        ChatMemberAdministrator,
+        ChatMemberMember,
+        ChatMemberRestricted,
+        ChatMemberBanned,
+        ChatMemberLeft,
+    ]:
         """https://core.telegram.org/bots/api/#getchatmember"""
         result = await self._send_request(
             "getChatMember",

diff --git a/tgram/types.py b/tgram/types.py
@@ -3652,31 +3652,32 @@ def _parse(
 
 
 class ChatMember(Type_):
-    def __init__(
-        self,
-        status: "str",
-        user: "User",
-        is_anonymous: "bool",
-        custom_title: "str" = None,
-        me: "tgram.TgBot" = None,
-        json: "dict" = None,
-    ):
-        super().__init__(me=me, json=json)
-        self.status = status
-        self.user = user
-        self.is_anonymous = is_anonymous
-        self.custom_title = custom_title
-
     @staticmethod
-    def _parse(me: "tgram.TgBot" = None, d: dict = None) -> Optional["ChatMember"]:
-        return (
-            ChatMember(
-                me=me,
-                json=d,
-                status=d.get("status"),
-                user=User._parse(me=me, d=d.get("user")),
-                is_anonymous=d.get("is_anonymous"),
-                custom_title=d.get("custom_title"),
+    def _parse(
+        me: "tgram.TgBot" = None, d: dict = None
+    ) -> Optional[
+        Union[
+            "ChatMemberOwner",
+            "ChatMemberAdministrator",
+            "ChatMemberMember",
+            "ChatMemberRestricted",
+            "ChatMemberBanned",
+            "ChatMemberLeft",
+        ]
+    ]:
+        return (
+            (
+                ChatMemberOwner._parse(me=me, d=d)
+                if d.get("status") == "creator"
+                else ChatMemberAdministrator._parse(me=me, d=d)
+                if d.get("status") == "administrator"
+                else ChatMemberMember._parse(me=me, d=d)
+                if d.get("status") == "member"
+                else ChatMemberRestricted._parse(me=me, d=d)
+                if d.get("status") == "restricted"
+                else ChatMemberLeft._parse(me=me, d=d)
+                if d.get("status") == "left"
+                else ChatMemberBanned._parse(me=me, d=d)
             )
             if d
             else None