Extend support for pyrogram types

pyrogram/client/__init__.py

@@ -20,3 +20,4 @@ from .chat_action import ChatAction
 from .client import Client
 from .parse_mode import ParseMode
 from .emoji import Emoji
+from . import utils

pyrogram/client/client.py

@@ -46,6 +46,7 @@ from pyrogram.api.errors import (
     PasswordHashInvalid, FloodWait, PeerIdInvalid, FilePartMissing,
     ChatAdminRequired, FirstnameInvalid, PhoneNumberBanned,
     VolumeLocNotFound, UserMigrate)
+from pyrogram.client import utils
 from pyrogram.crypto import AES
 from pyrogram.session import Auth, Session
 from pyrogram.session.internals import MsgId
@@ -704,77 +705,43 @@ class Client:
                 chats = {i.id: i for i in update[2]}
                 update = update[0]
 
-                if isinstance(update, types.UpdateNewMessage):
-                    message = update.message  # type: types.Message
+                valid_updates = (types.UpdateNewMessage, types.UpdateNewChannelMessage,
+                                 types.UpdateEditMessage, types.UpdateEditChannelMessage)
+
+                if isinstance(update, valid_updates):
+                    message = update.message
 
                     if isinstance(message, types.Message):
-                        from_user = users[message.from_id]  # type: types.User
-                        text = message.message or None
+                        m = utils.parse_message(message, users, chats)
 
-                        if isinstance(message.to_id, types.PeerUser):
-                            to_user_id = message.to_id.user_id
-                            to_user = users[to_user_id]  # type: types.User
+                        if message.reply_to_msg_id:
+                            rm = self.get_messages(m.chat.id, [message.reply_to_msg_id])
 
-                            chat = pyrogram.Chat(
-                                id=to_user_id,
-                                type="private",
-                                username=to_user.username,
-                                first_name=to_user.first_name,
-                                last_name=to_user.last_name
-                            )
-                        else:
-                            to_group_id = message.to_id.chat_id
-                            to_group = chats[to_group_id]  # type: types.Chat
+                            message = rm.messages[0]
 
-                            chat = pyrogram.Chat(
-                                id=-to_group_id,
-                                type="group",
-                                title=to_group.title,
-                                all_members_are_administrators=to_group.admins_enabled
-                            )
-                    else:
-                        continue
-                elif isinstance(update, types.UpdateNewChannelMessage):
-                    message = update.message  # type: types.Message
-
-                    if isinstance(message, types.Message):
-                        from_user = users.get(message.from_id, None)  # type: types.User or None
-                        to_channel_id = message.to_id.channel_id
-                        to_channel = chats[to_channel_id]  # type: types.Channel
-                        text = message.message or None
-
-                        chat = pyrogram.Chat(
-                            id=int("-100" + str(to_channel_id)),
-                            type="supergroup" if to_channel.megagroup else "channel",
-                            title=to_channel.title,
-                            username=to_channel.username
-                        )
-                    else:
-                        continue
+                            if isinstance(message, types.Message):
+                                m.reply_to_message = utils.parse_message(
+                                    message,
+                                    {i.id: i for i in rm.users},
+                                    {i.id: i for i in rm.chats}
+                                )
+                            else:
+                                continue
                 else:
                     continue
 
+                edit = isinstance(update, (types.UpdateEditMessage, types.UpdateEditChannelMessage))
+
                 u = pyrogram.Update(
                     update_id=0,
-                    message=pyrogram.Message(
-                        message_id=message.id,
-                        date=message.date,
-                        chat=chat,
-                        from_user=pyrogram.User(
-                            id=from_user.id,
-                            is_bot=from_user.bot,
-                            first_name=from_user.first_name,
-                            last_name=from_user.last_name,
-                            username=from_user.username,
-                            language_code=from_user.lang_code
-                        ) if from_user else None,
-                        text=text
-                    )
+                    message=(m if m.chat.type is not "channel" else None) if not edit else None,
+                    edited_message=(m if m.chat.type is not "channel" else None) if edit else None,
+                    channel_post=(m if m.chat.type is "channel" else None) if not edit else None,
+                    edited_channel_post=(m if m.chat.type is "channel" else None) if edit else None
                 )
 
                 if self.update_handler:
                     self.update_handler(self, u)
-                    # self.update_handler(self, update, users, chats)
             except Exception as e:
                 log.error(e, exc_info=True)
 

pyrogram/client/utils.py

@@ -0,0 +1,120 @@
+import pyrogram
+from pyrogram.api import types
+
+ENTITIES = {
+    types.MessageEntityMention.ID: "mention",
+    types.MessageEntityHashtag.ID: "hashtag",
+    types.MessageEntityBotCommand.ID: "bot_command",
+    types.MessageEntityUrl.ID: "url",
+    types.MessageEntityEmail.ID: "email",
+    types.MessageEntityBold.ID: "bold",
+    types.MessageEntityItalic.ID: "italic",
+    types.MessageEntityCode.ID: "code",
+    types.MessageEntityPre.ID: "pre",
+    types.MessageEntityTextUrl.ID: "text_link",
+}
+
+
+def parse_entities(entities: list):
+    output_entities = []
+
+    for entity in entities:
+        entity_type = ENTITIES.get(entity.ID, None)
+
+        if entity_type:
+            output_entities.append(pyrogram.MessageEntity(
+                type=entity_type,
+                offset=entity.offset,
+                length=entity.length,
+                url=getattr(entity, "url", None)
+            ))
+
+    return output_entities
+
+
+def parse_user(user: types.User):
+    return pyrogram.User(
+        id=user.id,
+        is_bot=user.bot,
+        first_name=user.first_name,
+        last_name=user.last_name,
+        username=user.username,
+        language_code=user.lang_code
+    ) if user else None
+
+
+def parse_user_chat(user: types.User):
+    return pyrogram.Chat(
+        id=user.id,
+        type="private",
+        username=user.username,
+        first_name=user.first_name,
+        last_name=user.last_name
+    )
+
+
+def parse_chat_chat(chat: types.Chat):
+    return pyrogram.Chat(
+        id=-chat.id,
+        type="group",
+        title=chat.title,
+        all_members_are_administrators=chat.admins_enabled
+    )
+
+
+def parse_channel_chat(channel: types.Channel):
+    return pyrogram.Chat(
+        id=int("-100" + str(channel.id)),
+        type="supergroup" if channel.megagroup else "channel",
+        title=channel.title,
+        username=channel.username
+    )
+
+
+def parse_message(message: types.Message, users: dict, chats: dict):
+    from_user = users.get(message.from_id, None)  # type: types.User
+
+    if isinstance(message.to_id, types.PeerUser):
+        chat = parse_user_chat(users[message.to_id.user_id])
+    elif isinstance(message.to_id, types.PeerChat):
+        chat = parse_chat_chat(chats[message.to_id.chat_id])
+    else:
+        chat = parse_channel_chat(chats[message.to_id.channel_id])
+
+    entities = parse_entities(message.entities)
+
+    forward_from = None
+    forward_from_chat = None
+    forward_from_message_id = None
+    forward_signature = None
+    forward_date = None
+
+    forward_header = message.fwd_from  # type: types.MessageFwdHeader
+
+    if forward_header:
+        forward_date = forward_header.date
+
+        if forward_header.from_id:
+            forward_from = parse_user(users[forward_header.from_id])
+        else:
+            forward_from_chat = parse_channel_chat(chats[forward_header.channel_id])
+            forward_from_message_id = forward_header.channel_post
+            forward_signature = forward_header.post_author
+
+    return pyrogram.Message(
+        message_id=message.id,
+        date=message.date,
+        chat=chat,
+        from_user=parse_user(from_user),
+        text=message.message or None if message.media is None else None,
+        caption=message.message or None if message.media is not None else None,
+        entities=entities or None if message.media is None else None,
+        caption_entities=entities or None if message.media is not None else None,
+        author_signature=message.post_author,
+        forward_from=forward_from,
+        forward_from_chat=forward_from_chat,
+        forward_from_message_id=forward_from_message_id,
+        forward_signature=forward_signature,
+        forward_date=forward_date,
+        edit_date=message.edit_date
+    )