diff --git a/tgram/types/_external_reply_info.py b/tgram/types/_external_reply_info.py index 4e7e51f..544c426 100644 --- a/tgram/types/_external_reply_info.py +++ b/tgram/types/_external_reply_info.py @@ -2,6 +2,7 @@ from .type_ import Type_ from typing import List, Optional +from tgram.utils import message_origin_parse class ExternalReplyInfo(Type_): @@ -147,19 +148,7 @@ def _parse( ExternalReplyInfo( me=me, json=d, - origin=None - if not d.get("origin") - else tgram.types.MessageOriginUser._parse(me=me, d=d.get("origin")) - if d["origin"].get("sender_user") - else tgram.types.MessageOriginHiddenUser._parse( - me=me, d=d.get("origin") - ) - if d["origin"].get("sender_user_name") - else tgram.types.MessageOriginChat._parse(me=me, d=d.get("origin")) - if d["origin"].get("sender_chat") - else tgram.types.MessageOriginChannel._parse(me=me, d=d.get("origin")) - if d["origin"].get("author_signature") - else None, + origin=message_origin_parse(d.get("origin"), me), chat=tgram.types.Chat._parse(me=me, d=d.get("chat")), message_id=d.get("message_id"), link_preview_options=tgram.types.LinkPreviewOptions._parse( diff --git a/tgram/types/_message.py b/tgram/types/_message.py index 8a1b199..6ace410 100644 --- a/tgram/types/_message.py +++ b/tgram/types/_message.py @@ -3,7 +3,7 @@ from typing import List, Optional from tgram import bound -from tgram.utils import String +from tgram.utils import String, message_origin_parse class Message(Type_, bound.MessageB): @@ -493,25 +493,7 @@ def _parse( me=me, d=d.get("sender_business_bot") ), business_connection_id=d.get("business_connection_id"), - forward_origin=None - if not d.get("forward_origin") - else tgram.types.MessageOriginUser._parse( - me=me, d=d.get("forward_origin") - ) - if d["forward_origin"].get("sender_user") - else tgram.types.MessageOriginHiddenUser._parse( - me=me, d=d.get("forward_origin") - ) - if d["forward_origin"].get("sender_user_name") - else tgram.types.MessageOriginChat._parse( - me=me, d=d.get("forward_origin") - ) - if d["forward_origin"].get("sender_chat") - else tgram.types.MessageOriginChannel._parse( - me=me, d=d.get("forward_origin") - ) - if d["forward_origin"].get("author_signature") - else None, + forward_origin=message_origin_parse(d.get("forward_origin"), me), is_topic_message=d.get("is_topic_message"), is_automatic_forward=d.get("is_automatic_forward"), reply_to_message=tgram.types.Message._parse( diff --git a/tgram/utils.py b/tgram/utils.py index 3d3667c..3dc1d58 100644 --- a/tgram/utils.py +++ b/tgram/utils.py @@ -9,7 +9,7 @@ import inspect from pathlib import Path -from typing import List, Union, TypedDict +from typing import List, Union, TypedDict, Optional from struct import unpack from io import BytesIO from json import dumps @@ -376,6 +376,25 @@ async def compose(bots: List["tgram.TgBot"]): return await asyncio.wait(tasks) +def message_origin_parse( + d: Optional[dict] = None, me: Optional["tgram.TgBot"] = None +) -> Optional["tgram.types.MessageOrigin"]: + if d is None: + return None + + origin_type = d["type"] + + return ( + tgram.types.MessageOriginUser._parse(me=me, d=d) + if origin_type == "user" + else tgram.types.MessageOriginHiddenUser._parse(me=me, d=d) + if origin_type == "hidden_user" + else tgram.types.MessageOriginChat._parse(me=me, d=d) + if origin_type == "chat" + else tgram.types.MessageOriginChannel._parse(me=me, d=d) + ) + + def b64_decode(s: str) -> bytes: return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))