Refactor Sticker parsing

delivrance · delivrance · commit 2e4651401216 · 2022-08-14T11:19:01.000+02:00
diff --git a/pyrogram/methods/messages/get_custom_emoji_stickers.py b/pyrogram/methods/messages/get_custom_emoji_stickers.py
@@ -47,15 +47,7 @@ async def get_custom_emoji_stickers(
         stickers = []
         for item in result:
             attributes = {type(i): i for i in item.attributes}
-
-            sticker = await types.Sticker._parse(
-                self, item,
-                attributes[raw.types.DocumentAttributeImageSize] if raw.types.DocumentAttributeImageSize in attributes else None,
-                attributes[raw.types.DocumentAttributeCustomEmoji],
-                attributes[raw.types.DocumentAttributeFilename].file_name,
-                attributes[raw.types.DocumentAttributeVideo] if raw.types.DocumentAttributeVideo in attributes else None
-            )
-
+            sticker = await types.Sticker._parse(self, item, attributes)
             stickers.append(sticker)
 
         return pyrogram.types.List(stickers)
diff --git a/pyrogram/types/messages_and_media/message.py b/pyrogram/types/messages_and_media/message.py
@@ -696,13 +696,7 @@ async def _parse(
                             animation = types.Animation._parse(client, doc, video_attributes, file_name)
                             media_type = enums.MessageMediaType.ANIMATION
                         elif raw.types.DocumentAttributeSticker in attributes:
-                            sticker = await types.Sticker._parse(
-                                client, doc,
-                                attributes.get(raw.types.DocumentAttributeImageSize, None),
-                                attributes[raw.types.DocumentAttributeSticker],
-                                file_name,
-                                attributes.get(raw.types.DocumentAttributeVideo, None)
-                            )
+                            sticker = await types.Sticker._parse(client, doc, attributes)
                             media_type = enums.MessageMediaType.STICKER
                         elif raw.types.DocumentAttributeVideo in attributes:
                             video_attributes = attributes[raw.types.DocumentAttributeVideo]
diff --git a/pyrogram/types/messages_and_media/sticker.py b/pyrogram/types/messages_and_media/sticker.py
@@ -17,7 +17,7 @@
 #  along with Pyrogram.  If not, see <http://www.gnu.org/licenses/>.
 
 from datetime import datetime
-from typing import List
+from typing import List, Dict, Type
 
 import pyrogram
 from pyrogram import raw, utils
@@ -50,9 +50,6 @@ class Sticker(Object):
         is_video (``bool``):
             True, if the sticker is a video sticker
 
-        duration (``int``):
-            Video sticker duration in seconds.
-
         file_name (``str``, *optional*):
             Sticker file name.
 
@@ -87,7 +84,6 @@ def __init__(
         height: int,
         is_animated: bool,
         is_video: bool,
-        duration: int = None,
         file_name: str = None,
         mime_type: str = None,
         file_size: int = None,
@@ -150,11 +146,16 @@ async def _get_sticker_set_name(invoke, input_sticker_set_id):
     async def _parse(
         client,
         sticker: "raw.types.Document",
-        image_size_attributes: "raw.types.DocumentAttributeImageSize",
-        sticker_attributes: "raw.types.DocumentAttributeSticker",
-        file_name: str,
-        video_attributes: "raw.types.DocumentAttributeVideo"
+        document_attributes: Dict[Type["raw.base.DocumentAttribute"], "raw.base.DocumentAttribute"],
     ) -> "Sticker":
+        sticker_attributes = document_attributes.get(
+            raw.types.DocumentAttributeSticker,
+            document_attributes[raw.types.DocumentAttributeCustomEmoji]
+        )
+        image_size_attributes = document_attributes.get(raw.types.DocumentAttributeImageSize, None)
+        file_name = getattr(document_attributes.get(raw.types.DocumentAttributeFilename, None), "file_name", None)
+        video_attributes = document_attributes.get(raw.types.DocumentAttributeVideo, None)
+
         sticker_set = sticker_attributes.stickerset
 
         if isinstance(sticker_set, raw.types.InputStickerSetID):
@@ -175,11 +176,22 @@ async def _parse(
                 file_unique_type=FileUniqueType.DOCUMENT,
                 media_id=sticker.id
             ).encode(),
-            width=image_size_attributes.w if image_size_attributes else (video_attributes.w if video_attributes else 512),
-            height=image_size_attributes.h if image_size_attributes else (video_attributes.h if video_attributes else 512),
+            width=(
+                image_size_attributes.w
+                if image_size_attributes
+                else video_attributes.w
+                if video_attributes
+                else 512
+            ),
+            height=(
+                image_size_attributes.h
+                if image_size_attributes
+                else video_attributes.h
+                if video_attributes
+                else 512
+            ),
             is_animated=sticker.mime_type == "application/x-tgsticker",
             is_video=sticker.mime_type == "video/webm",
-            duration=video_attributes.duration if video_attributes else None,
             # TODO: mask_position
             set_name=set_name,
             emoji=sticker_attributes.alt or None,
