Refactor Story parsing

pyrogram/types/messages_and_media/story.py

@@ -167,8 +167,7 @@ class Story(Object, Update):
         if isinstance(stories, raw.types.StoryItemDeleted):
             return await types.StoryDeleted._parse(client, stories, users, chats, peer)
 
-        entities = [types.MessageEntity._parse(client, entity, {}) for entity in stories.entities]
+        entities = [e for e in (types.MessageEntity._parse(client, entity, {}) for entity in stories.entities) if e]
-        entities = types.List(filter(lambda x: x is not None, entities)) or None
 
         photo = None
         video = None
@@ -179,26 +178,26 @@ class Story(Object, Update):
         allowed_users = None
         denied_chats = None
         denied_users = None
+        media_type = None
 
-        if stories.media:
+        if isinstance(stories.media, raw.types.MessageMediaPhoto):
-            media_type = None
+            photo = types.Photo._parse(client, stories.media.photo, stories.media.ttl_seconds)
-            if isinstance(stories.media, raw.types.MessageMediaPhoto):
+            media_type = enums.MessageMediaType.PHOTO
-                photo = types.Photo._parse(client, stories.media.photo, stories.media.ttl_seconds)
+        else:
-                media_type = enums.MessageMediaType.PHOTO
+            doc = stories.media.document
-            else:
+            attributes = {type(i): i for i in doc.attributes}
-                doc = stories.media.document
+            video_attributes = attributes.get(raw.types.DocumentAttributeVideo, None)
-                attributes = {type(i): i for i in doc.attributes}
+            video = types.Video._parse(client, doc, video_attributes, None)
-                video_attributes = attributes.get(raw.types.DocumentAttributeVideo, None)
+            media_type = enums.MessageMediaType.VIDEO
-                video = types.Video._parse(client, doc, video_attributes, None)
-                media_type = enums.MessageMediaType.VIDEO
 
-        peer_id = None
+        if isinstance(peer, raw.types.InputPeerSelf):
-        if hasattr(peer, "user_id"):
+            r = await client.invoke(raw.functions.users.GetUsers(id=[raw.types.InputPeerSelf()]))
-            peer_id = peer.user_id
+            peer_id = r[0].id
-        elif hasattr(peer, "chat_id"):
+            users.update({i.id: i for i in r})
-            peer_id = peer.chat_id
+        elif isinstance(peer, (raw.types.InputPeerUser, raw.types.InputPeerChannel)):
-        elif hasattr(peer, "channel_id"):
+            peer_id = utils.get_input_peer_id(peer)
-            peer_id = peer.channel_id
+        else:
+            peer_id = utils.get_raw_peer_id(peer)
 
         if isinstance(peer, (raw.types.PeerUser, raw.types.InputPeerUser)) and peer_id not in users:
             try:
@@ -213,31 +212,20 @@ class Story(Object, Update):
                 pass
             else:
                 users.update({i.id: i for i in r})
-        if isinstance(peer, raw.types.InputPeerSelf):
-            r = await client.invoke(
-                raw.functions.users.GetUsers(
-                    id=[
-                        raw.types.InputPeerSelf()
-                    ]
-                )
-            )
-            peer_id = r[0].id
-            users.update({i.id: i for i in r})
 
         from_user = types.User._parse(client, users.get(peer_id, None))
         sender_chat = types.Chat._parse_channel_chat(client, chats[peer_id]) if not from_user else None
 
+        privacy_map = {
+            raw.types.PrivacyValueAllowAll: enums.StoriesPrivacyRules.PUBLIC,
+            raw.types.PrivacyValueAllowCloseFriends: enums.StoriesPrivacyRules.CLOSE_FRIENDS,
+            raw.types.PrivacyValueAllowContacts: enums.StoriesPrivacyRules.CONTACTS,
+            raw.types.PrivacyValueDisallowAll: enums.StoriesPrivacyRules.PRIVATE,
+            raw.types.PrivacyValueDisallowContacts: enums.StoriesPrivacyRules.NO_CONTACTS
+        }
+
         for priv in stories.privacy:
-            if isinstance(priv, raw.types.PrivacyValueAllowAll):
+            privacy = privacy_map.get(type(priv), None)
-                privacy = enums.StoriesPrivacyRules.PUBLIC
-            elif isinstance(priv, raw.types.PrivacyValueAllowCloseFriends):
-                privacy = enums.StoriesPrivacyRules.CLOSE_FRIENDS
-            elif isinstance(priv, raw.types.PrivacyValueAllowContacts):
-                privacy = enums.StoriesPrivacyRules.CONTACTS
-            elif isinstance(priv, raw.types.PrivacyValueDisallowAll):
-                privacy = enums.StoriesPrivacyRules.PRIVATE
-            elif isinstance(priv, raw.types.PrivacyValueDisallowContacts):
-                privacy = enums.StoriesPrivacyRules.NO_CONTACTS
 
             if isinstance(priv, raw.types.PrivacyValueAllowUsers):
                 allowed_users = priv.users
@@ -261,8 +249,8 @@ class Story(Object, Update):
             contacts=stories.contacts,
             selected_contacts=stories.selected_contacts,
             caption=stories.caption,
-            caption_entities=entities,
+            caption_entities=entities or None,
-            views=types.StoryViews._parse(client, stories.views),
+            views=types.StoryViews._parse(client, stories.views) if stories.views else None,
             privacy=privacy,
             allowed_chats=allowed_chats,
             denied_chats=denied_chats,

pyrogram/types/messages_and_media/story_deleted.py

@@ -18,7 +18,7 @@
 
 import pyrogram
 
-from pyrogram import raw, types
+from pyrogram import raw, types, utils
 from pyrogram.errors import PeerIdInvalid
 from typing import Union
 from ..object import Object
@@ -63,13 +63,14 @@ class StoryDeleted(Object, Update):
         from_user = None
         sender_chat = None
 
-        peer_id = None
+        if isinstance(peer, raw.types.InputPeerSelf):
-        if hasattr(peer, "user_id"):
+            r = await client.invoke(raw.functions.users.GetUsers(id=[raw.types.InputPeerSelf()]))
-            peer_id = peer.user_id
+            peer_id = r[0].id
-        elif hasattr(peer, "chat_id"):
+            users.update({i.id: i for i in r})
-            peer_id = peer.chat_id
+        elif isinstance(peer, (raw.types.InputPeerUser, raw.types.InputPeerChannel)):
-        elif hasattr(peer, "channel_id"):
+            peer_id = utils.get_input_peer_id(peer)
-            peer_id = peer.channel_id
+        else:
+            peer_id = utils.get_raw_peer_id(peer)
 
         if isinstance(peer, (raw.types.PeerUser, raw.types.InputPeerUser)) and peer_id not in users:
             try:
@@ -84,16 +85,6 @@ class StoryDeleted(Object, Update):
                 pass
             else:
                 users.update({i.id: i for i in r})
-        if isinstance(peer, raw.types.InputPeerSelf):
-            r = await client.invoke(
-                raw.functions.users.GetUsers(
-                    id=[
-                        raw.types.InputPeerSelf()
-                    ]
-                )
-            )
-            peer_id = r[0].id
-            users.update({i.id: i for i in r})
 
         from_user = types.User._parse(client, users.get(peer_id, None))
         sender_chat = types.Chat._parse_channel_chat(client, chats[peer_id]) if not from_user else None

pyrogram/types/messages_and_media/story_skipped.py

@@ -79,13 +79,14 @@ class StorySkipped(Object, Update):
         from_user = None
         sender_chat = None
 
-        peer_id = None
+        if isinstance(peer, raw.types.InputPeerSelf):
-        if hasattr(peer, "user_id"):
+            r = await client.invoke(raw.functions.users.GetUsers(id=[raw.types.InputPeerSelf()]))
-            peer_id = peer.user_id
+            peer_id = r[0].id
-        elif hasattr(peer, "chat_id"):
+            users.update({i.id: i for i in r})
-            peer_id = peer.chat_id
+        elif isinstance(peer, (raw.types.InputPeerUser, raw.types.InputPeerChannel)):
-        elif hasattr(peer, "channel_id"):
+            peer_id = utils.get_input_peer_id(peer)
-            peer_id = peer.channel_id
+        else:
+            peer_id = utils.get_raw_peer_id(peer)
 
         if isinstance(peer, (raw.types.PeerUser, raw.types.InputPeerUser)) and peer_id not in users:
             try:
@@ -100,16 +101,6 @@ class StorySkipped(Object, Update):
                 pass
             else:
                 users.update({i.id: i for i in r})
-        if isinstance(peer, raw.types.InputPeerSelf):
-            r = await client.invoke(
-                raw.functions.users.GetUsers(
-                    id=[
-                        raw.types.InputPeerSelf()
-                    ]
-                )
-            )
-            peer_id = r[0].id
-            users.update({i.id: i for i in r})
 
         from_user = types.User._parse(client, users.get(peer_id, None))
         sender_chat = types.Chat._parse_channel_chat(client, chats[peer_id]) if not from_user else None

pyrogram/types/messages_and_media/story_views.py

@@ -71,5 +71,5 @@ class StoryViews(Object):
                 types.Reaction._parse_count(client, reaction)
                 for reaction in getattr(storyviews, "reactions", [])
             ] or None,
-            recent_viewers=getattr(storyviews, "recent_viewers", None),
+            recent_viewers=getattr(storyviews, "recent_viewers", None) or None,
         )

pyrogram/utils.py

@@ -245,6 +245,19 @@ def get_raw_peer_id(peer: raw.base.Peer) -> Optional[int]:
 
     return None
 
+def get_input_peer_id(peer: raw.base.InputPeer) -> Optional[int]:
+    """Get the raw peer id from a InputPeer object"""
+    if isinstance(peer, raw.types.InputPeerUser):
+        return peer.user_id
+
+    if isinstance(peer, raw.types.InputPeerChat):
+        return peer.chat_id
+
+    if isinstance(peer, raw.types.InputPeerChannel):
+        return peer.channel_id
+
+    return None
+
 
 def get_peer_id(peer: raw.base.Peer) -> int:
     """Get the non-raw peer id from a Peer object"""