Return custom.Message from the TelegramClient

This commit is contained in:
Lonami Exo 2018-05-31 12:52:03 +02:00
parent 5aed494aac
commit b241d80958

View File

@ -89,10 +89,11 @@ from .tl.types import (
MessageMediaWebPage, ChannelParticipantsSearch, PhotoSize, PhotoCachedSize, MessageMediaWebPage, ChannelParticipantsSearch, PhotoSize, PhotoCachedSize,
PhotoSizeEmpty, MessageService, ChatParticipants, User, WebPage, PhotoSizeEmpty, MessageService, ChatParticipants, User, WebPage,
ChannelParticipantsBanned, ChannelParticipantsKicked, ChannelParticipantsBanned, ChannelParticipantsKicked,
InputMessagesFilterEmpty InputMessagesFilterEmpty, UpdatesCombined
) )
from .tl.types.messages import DialogsSlice from .tl.types.messages import DialogsSlice
from .tl.types.account import PasswordInputSettings, NoPassword from .tl.types.account import PasswordInputSettings, NoPassword
from .tl import custom
from .extensions import markdown, html from .extensions import markdown, html
__log__ = logging.getLogger(__name__) __log__ = logging.getLogger(__name__)
@ -599,9 +600,10 @@ class TelegramClient(TelegramBareClient):
if _total: if _total:
_total[0] = getattr(r, 'count', len(r.dialogs)) _total[0] = getattr(r, 'count', len(r.dialogs))
messages = {m.id: m for m in r.messages}
entities = {utils.get_peer_id(x): x entities = {utils.get_peer_id(x): x
for x in itertools.chain(r.users, r.chats)} for x in itertools.chain(r.users, r.chats)}
messages = {m.id: custom.Message(self, m, entities, None)
for m in r.messages}
# Happens when there are pinned dialogs # Happens when there are pinned dialogs
if len(r.dialogs) > limit: if len(r.dialogs) > limit:
@ -652,8 +654,7 @@ class TelegramClient(TelegramBareClient):
""" """
return list(self.iter_drafts()) return list(self.iter_drafts())
@staticmethod def _get_response_message(self, request, result, input_chat):
def _get_response_message(request, result):
""" """
Extracts the response message known a request and Update result. Extracts the response message known a request and Update result.
The request may also be the ID of the message to match. The request may also be the ID of the message to match.
@ -672,26 +673,36 @@ class TelegramClient(TelegramBareClient):
if isinstance(result, UpdateShort): if isinstance(result, UpdateShort):
updates = [result.update] updates = [result.update]
elif isinstance(result, Updates): entities = {}
elif isinstance(result, (Updates, UpdatesCombined)):
updates = result.updates updates = result.updates
entities = {utils.get_peer_id(x): x
for x in itertools.chain(result.users, result.chats)}
else: else:
return return
found = None
for update in updates: for update in updates:
if isinstance(update, (UpdateNewChannelMessage, UpdateNewMessage)): if isinstance(update, (UpdateNewChannelMessage, UpdateNewMessage)):
if update.message.id == msg_id: if update.message.id == msg_id:
return update.message found = update.message
break
elif (isinstance(update, UpdateEditMessage) and elif (isinstance(update, UpdateEditMessage) and
not isinstance(request.peer, InputPeerChannel)): not isinstance(request.peer, InputPeerChannel)):
if request.id == update.message.id: if request.id == update.message.id:
return update.message found = update.message
break
elif (isinstance(update, UpdateEditChannelMessage) and elif (isinstance(update, UpdateEditChannelMessage) and
utils.get_peer_id(request.peer) == utils.get_peer_id(request.peer) ==
utils.get_peer_id(update.message.to_id)): utils.get_peer_id(update.message.to_id)):
if request.id == update.message.id: if request.id == update.message.id:
return update.message found = update.message
break
if found:
return custom.Message(self, found, entities, input_chat)
def _parse_message_text(self, message, parse_mode): def _parse_message_text(self, message, parse_mode):
""" """
@ -839,17 +850,18 @@ class TelegramClient(TelegramBareClient):
result = self(request) result = self(request)
if isinstance(result, UpdateShortSentMessage): if isinstance(result, UpdateShortSentMessage):
return Message( to_id, cls = utils.resolve_id(utils.get_peer_id(entity))
return custom.Message(self, Message(
id=result.id, id=result.id,
to_id=entity, to_id=cls(to_id),
message=message, message=message,
date=result.date, date=result.date,
out=result.out, out=result.out,
media=result.media, media=result.media,
entities=result.entities entities=result.entities
) ), {}, input_chat=entity)
return self._get_response_message(request, result) return self._get_response_message(request, result, entity)
def forward_messages(self, entity, messages, from_peer=None): def forward_messages(self, entity, messages, from_peer=None):
""" """
@ -895,13 +907,20 @@ class TelegramClient(TelegramBareClient):
to_peer=entity to_peer=entity
) )
result = self(req) result = self(req)
if isinstance(result, (Updates, UpdatesCombined)):
entities = {utils.get_peer_id(x): x
for x in itertools.chain(result.users, result.chats)}
else:
entities = {}
random_to_id = {} random_to_id = {}
id_to_message = {} id_to_message = {}
for update in result.updates: for update in result.updates:
if isinstance(update, UpdateMessageID): if isinstance(update, UpdateMessageID):
random_to_id[update.random_id] = update.id random_to_id[update.random_id] = update.id
elif isinstance(update, (UpdateNewMessage, UpdateNewChannelMessage)): elif isinstance(update, (UpdateNewMessage, UpdateNewChannelMessage)):
id_to_message[update.message.id] = update.message id_to_message[update.message.id] = custom.Message(
self, update.message, entities, input_chat=entity)
result = [id_to_message[random_to_id[rnd]] for rnd in req.random_id] result = [id_to_message[random_to_id[rnd]] for rnd in req.random_id]
return result[0] if single else result return result[0] if single else result
@ -962,16 +981,16 @@ class TelegramClient(TelegramBareClient):
message = entity message = entity
entity = entity.to_id entity = entity.to_id
entity = self.get_input_entity(entity)
text, msg_entities = self._parse_message_text(text, parse_mode) text, msg_entities = self._parse_message_text(text, parse_mode)
request = EditMessageRequest( request = EditMessageRequest(
peer=self.get_input_entity(entity), peer=entity,
id=self._get_message_id(message), id=self._get_message_id(message),
message=text, message=text,
no_webpage=not link_preview, no_webpage=not link_preview,
entities=msg_entities entities=msg_entities
) )
result = self(request) return self._get_response_message(request, self(request), entity)
return self._get_response_message(request, result)
def delete_messages(self, entity, message_ids, revoke=True): def delete_messages(self, entity, message_ids, revoke=True):
""" """
@ -1189,8 +1208,7 @@ class TelegramClient(TelegramBareClient):
# IDs are returned in descending order. # IDs are returned in descending order.
last_id = message.id last_id = message.id
self._make_message_friendly(message, entities) yield custom.Message(self, message, entities, entity)
yield message
have += 1 have += 1
if len(r.messages) < request.limit: if len(r.messages) < request.limit:
@ -1204,34 +1222,6 @@ class TelegramClient(TelegramBareClient):
time.sleep(max(wait_time - (time.time() - start), 0)) time.sleep(max(wait_time - (time.time() - start), 0))
@staticmethod
def _make_message_friendly(message, entities):
"""
Add a few extra attributes to the :tl:`Message` to be friendlier.
To make messages more friendly, always add message
to service messages, and action to normal messages.
"""
# TODO Create an actual friendlier class
message.message = getattr(message, 'message', None)
message.action = getattr(message, 'action', None)
message.to = entities[utils.get_peer_id(message.to_id)]
message.sender = (
None if not message.from_id else
entities[utils.get_peer_id(message.from_id)]
)
if getattr(message, 'fwd_from', None):
message.fwd_from.sender = (
None if not message.fwd_from.from_id else
entities[utils.get_peer_id(message.fwd_from.from_id)]
)
message.fwd_from.channel = (
None if not message.fwd_from.channel_id else
entities[utils.get_peer_id(
PeerChannel(message.fwd_from.channel_id)
)]
)
def _iter_ids(self, entity, ids, total): def _iter_ids(self, entity, ids, total):
""" """
Special case for `iter_messages` when it should only fetch some IDs. Special case for `iter_messages` when it should only fetch some IDs.
@ -1253,8 +1243,7 @@ class TelegramClient(TelegramBareClient):
if isinstance(message, MessageEmpty): if isinstance(message, MessageEmpty):
yield None yield None
else: else:
self._make_message_friendly(message, entities) yield custom.Message(self, message, entities, entity)
yield message
def get_messages(self, *args, **kwargs): def get_messages(self, *args, **kwargs):
""" """
@ -1355,6 +1344,9 @@ class TelegramClient(TelegramBareClient):
if isinstance(message, int): if isinstance(message, int):
return message return message
if isinstance(message, custom.Message):
return message.original_message.id
try: try:
if message.SUBCLASS_OF_ID == 0x790009e3: if message.SUBCLASS_OF_ID == 0x790009e3:
# hex(crc32(b'Message')) = 0x790009e3 # hex(crc32(b'Message')) = 0x790009e3
@ -1676,7 +1668,8 @@ class TelegramClient(TelegramBareClient):
reply_to_msg_id=reply_to, reply_to_msg_id=reply_to,
message=caption, message=caption,
entities=msg_entities) entities=msg_entities)
return self._get_response_message(request, self(request)) return self._get_response_message(request, self(request),
entity)
as_image = utils.is_image(file) and not force_document as_image = utils.is_image(file) and not force_document
use_cache = InputPhoto if as_image else InputDocument use_cache = InputPhoto if as_image else InputDocument
@ -1774,7 +1767,7 @@ class TelegramClient(TelegramBareClient):
# send the media message to the desired entity. # send the media message to the desired entity.
request = SendMediaRequest(entity, media, reply_to_msg_id=reply_to, request = SendMediaRequest(entity, media, reply_to_msg_id=reply_to,
message=caption, entities=msg_entities) message=caption, entities=msg_entities)
msg = self._get_response_message(request, self(request)) msg = self._get_response_message(request, self(request), entity)
if msg and isinstance(file_handle, InputSizedFile): if msg and isinstance(file_handle, InputSizedFile):
# There was a response message and we didn't use cached # There was a response message and we didn't use cached
# version, so cache whatever we just sent to the database. # version, so cache whatever we just sent to the database.
@ -1840,7 +1833,7 @@ class TelegramClient(TelegramBareClient):
entity, reply_to_msg_id=reply_to, multi_media=media entity, reply_to_msg_id=reply_to, multi_media=media
)) ))
return [ return [
self._get_response_message(update.id, result) self._get_response_message(update.id, result, entity)
for update in result.updates for update in result.updates
if isinstance(update, UpdateMessageID) if isinstance(update, UpdateMessageID)
] ]