diff --git a/pyrogram/client/__init__.py b/pyrogram/client/__init__.py index abda4464..49214797 100644 --- a/pyrogram/client/__init__.py +++ b/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 diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py index 5f51161d..cb585566 100644 --- a/pyrogram/client/client.py +++ b/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) diff --git a/pyrogram/client/utils.py b/pyrogram/client/utils.py new file mode 100644 index 00000000..5d5b37ba --- /dev/null +++ b/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 + )