Refactor ChatColor

KurimuzonAkuma · Dec 6, 2023 · 7b199bd · 7b199bd
1 parent f85fdc8
commit 7b199bd
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 10 deletions.
diff --git a/pyrogram/types/user_and_chats/chat.py b/pyrogram/types/user_and_chats/chat.py
@@ -263,7 +263,7 @@ def _parse_user_chat(client, user: raw.types.User) -> "Chat":
             restrictions=types.List([types.Restriction._parse(r) for r in user.restriction_reason]) or None,
             dc_id=getattr(getattr(user, "photo", None), "dc_id", None),
             reply_color=types.ChatColor._parse(getattr(user, "color", None)),
-            profile_color=types.ChatColor._parse(getattr(user, "profile_color", None), for_profile=True),
+            profile_color=types.ChatColor._parse_profile_color(getattr(user, "profile_color", None)),
             client=client
         )
 

diff --git a/pyrogram/types/user_and_chats/chat_color.py b/pyrogram/types/user_and_chats/chat_color.py
@@ -44,19 +44,21 @@ def __init__(
         self.background_emoji_id = background_emoji_id
 
     @staticmethod
-    def _parse(color: "raw.types.PeerColor" = None, for_profile: bool = None) -> Optional["ChatColor"]:
+    def _parse(color: "raw.types.PeerColor" = None) -> Optional["ChatColor"]:
         if not color:
             return None
 
-        chat_color = getattr(color, "color", None)
+        return ChatColor(
+            color=enums.ReplyColor(color.color) if getattr(color, "color", None) else None,
+            background_emoji_id=getattr(color, "background_emoji_id", None)
+        )
 
-        if not chat_color is None:
-            if for_profile:
-                chat_color = enums.ProfileColor(color.color)
-            else:
-                chat_color = enums.ReplyColor(color.color)
+    @staticmethod
+    def _parse_profile_color(color: "raw.types.PeerColor" = None) -> Optional["ChatColor"]:
+        if not color:
+            return None
 
         return ChatColor(
-            color=chat_color,
+            color=enums.ProfileColor(color.color) if getattr(color, "color", None) else None,
             background_emoji_id=getattr(color, "background_emoji_id", None)
         )
diff --git a/pyrogram/types/user_and_chats/user.py b/pyrogram/types/user_and_chats/user.py
@@ -281,7 +281,7 @@ def _parse(client, user: "raw.base.User") -> Optional["User"]:
             photo=types.ChatPhoto._parse(client, user.photo, user.id, user.access_hash),
             restrictions=types.List([types.Restriction._parse(r) for r in user.restriction_reason]) or None,
             reply_color=types.ChatColor._parse(getattr(user, "color", None)),
-            profile_color=types.ChatColor._parse(getattr(user, "profile_color", None), for_profile=True),
+            profile_color=types.ChatColor._parse_profile_color(getattr(user, "profile_color", None)),
             client=client
         )