Reuse code for _get_entity_pair

Less error-prone, improved the function flow for all callers,
and removed some duplicate work.
This commit is contained in:
Lonami Exo
2019-05-01 17:52:32 +02:00
parent 6d004601d0
commit 3a1496c205
9 changed files with 46 additions and 64 deletions

View File

@@ -29,25 +29,16 @@ class Forward(ChatGetter, SenderGetter):
self.original_fwd = original
self._sender_id = original.from_id
self._sender = entities.get(original.from_id)
try:
self._input_sender =\
utils.get_input_peer(self._sender) if self._sender else None
except TypeError:
self._input_sender = None
self._sender, self._input_sender = utils._get_entity_pair(
original.from_id, entities, client._entity_cache)
self._broadcast = None
if original.channel_id:
self._chat_peer = types.PeerChannel(original.channel_id)
self._chat = entities.get(utils.get_peer_id(self._chat_peer))
else:
self._chat_peer = None
self._chat = None
try:
self._input_chat = \
utils.get_input_peer(self._chat) if self._chat else None
except TypeError:
self._input_chat = None
self._chat, self._input_chat = utils._get_entity_pair(
self.chat_id, entities, client._entity_cache)
else:
self._chat = self._input_chat = self._chat_peer = None
# TODO We could reload the message

View File

@@ -210,27 +210,20 @@ class Message(ChatGetter, SenderGetter, TLObject, abc.ABC):
known entities.
"""
self._client = client
self._sender = entities.get(self._sender_id)
if self._sender:
try:
self._input_sender = utils.get_input_peer(self._sender)
except TypeError:
self._input_sender = None
cache = client._entity_cache
self._chat = entities.get(self.chat_id)
self._input_chat = input_chat
if not self._input_chat and self._chat:
try:
self._input_chat = utils.get_input_peer(self._chat)
except TypeError:
self._input_chat = None
self._sender, self._input_sender = utils._get_entity_pair(
self.sender_id, entities, cache)
self._via_bot = entities.get(self.via_bot_id)
if self._via_bot:
try:
self._via_input_bot = utils.get_input_peer(self._via_bot)
except TypeError:
self._via_input_bot = None
self._chat, self._input_chat = utils._get_entity_pair(
self.chat_id, entities, cache)
if input_chat: # This has priority
self._input_chat = input_chat
if self.via_bot_id:
self._via_bot, self._via_input_bot = utils._get_entity_pair(
self.via_bot_id, entities, cache)
if self.fwd_from:
self._forward = Forward(self._client, self.fwd_from, entities)