Refactor Story parsing

This commit is contained in:
KurimuzonAkuma 2023-10-23 23:11:23 +03:00
parent bd5ea1407c
commit bdce1e7a9c
5 changed files with 61 additions and 78 deletions

View File

@ -167,8 +167,7 @@ class Story(Object, Update):
if isinstance(stories, raw.types.StoryItemDeleted): if isinstance(stories, raw.types.StoryItemDeleted):
return await types.StoryDeleted._parse(client, stories, users, chats, peer) 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 photo = None
video = None video = None
@ -179,9 +178,8 @@ class Story(Object, Update):
allowed_users = None allowed_users = None
denied_chats = None denied_chats = None
denied_users = None denied_users = None
if stories.media:
media_type = None media_type = None
if isinstance(stories.media, raw.types.MessageMediaPhoto): if isinstance(stories.media, raw.types.MessageMediaPhoto):
photo = types.Photo._parse(client, stories.media.photo, stories.media.ttl_seconds) photo = types.Photo._parse(client, stories.media.photo, stories.media.ttl_seconds)
media_type = enums.MessageMediaType.PHOTO media_type = enums.MessageMediaType.PHOTO
@ -192,13 +190,14 @@ class Story(Object, Update):
video = types.Video._parse(client, doc, video_attributes, None) video = types.Video._parse(client, doc, video_attributes, None)
media_type = enums.MessageMediaType.VIDEO 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: if isinstance(peer, (raw.types.PeerUser, raw.types.InputPeerUser)) and peer_id not in users:
try: try:
@ -213,31 +212,20 @@ class Story(Object, Update):
pass pass
else: else:
users.update({i.id: i for i in r}) 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)) 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 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: 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): if isinstance(priv, raw.types.PrivacyValueAllowUsers):
allowed_users = priv.users allowed_users = priv.users
@ -261,8 +249,8 @@ class Story(Object, Update):
contacts=stories.contacts, contacts=stories.contacts,
selected_contacts=stories.selected_contacts, selected_contacts=stories.selected_contacts,
caption=stories.caption, 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, privacy=privacy,
allowed_chats=allowed_chats, allowed_chats=allowed_chats,
denied_chats=denied_chats, denied_chats=denied_chats,

View File

@ -18,7 +18,7 @@
import pyrogram import pyrogram
from pyrogram import raw, types from pyrogram import raw, types, utils
from pyrogram.errors import PeerIdInvalid from pyrogram.errors import PeerIdInvalid
from typing import Union from typing import Union
from ..object import Object from ..object import Object
@ -63,13 +63,14 @@ class StoryDeleted(Object, Update):
from_user = None from_user = None
sender_chat = 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: if isinstance(peer, (raw.types.PeerUser, raw.types.InputPeerUser)) and peer_id not in users:
try: try:
@ -84,16 +85,6 @@ class StoryDeleted(Object, Update):
pass pass
else: else:
users.update({i.id: i for i in r}) 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)) 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 sender_chat = types.Chat._parse_channel_chat(client, chats[peer_id]) if not from_user else None

View File

@ -79,13 +79,14 @@ class StorySkipped(Object, Update):
from_user = None from_user = None
sender_chat = 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: if isinstance(peer, (raw.types.PeerUser, raw.types.InputPeerUser)) and peer_id not in users:
try: try:
@ -100,16 +101,6 @@ class StorySkipped(Object, Update):
pass pass
else: else:
users.update({i.id: i for i in r}) 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)) 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 sender_chat = types.Chat._parse_channel_chat(client, chats[peer_id]) if not from_user else None

View File

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

View File

@ -245,6 +245,19 @@ def get_raw_peer_id(peer: raw.base.Peer) -> Optional[int]:
return None 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: def get_peer_id(peer: raw.base.Peer) -> int:
"""Get the non-raw peer id from a Peer object""" """Get the non-raw peer id from a Peer object"""