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

@ -105,7 +105,8 @@ class CallbackQuery(EventBuilder):
def _set_client(self, client):
super()._set_client(client)
self._sender, self._input_sender = self._get_entity_pair(self.sender_id)
self._sender, self._input_sender = utils._get_entity_pair(
self.sender_id, self._entities, client._entity_cache)
@property
def id(self):

View File

@ -147,25 +147,11 @@ class EventCommon(ChatGetter, abc.ABC):
"""
self._client = client
if self._chat_peer:
self._chat, self._input_chat = self._get_entity_pair(self.chat_id)
self._chat, self._input_chat = utils._get_entity_pair(
self.chat_id, self._entities, client._entity_cache)
else:
self._chat = self._input_chat = None
def _get_entity_pair(self, entity_id):
"""
Returns ``(entity, input_entity)`` for the given entity ID.
"""
entity = self._entities.get(entity_id)
try:
input_entity = utils.get_input_peer(entity)
except TypeError:
try:
input_entity = self._client._entity_cache[entity_id]
except KeyError:
input_entity = None
return entity, input_entity
@property
def client(self):
"""

View File

@ -89,7 +89,8 @@ class InlineQuery(EventBuilder):
def _set_client(self, client):
super()._set_client(client)
self._sender, self._input_sender = self._get_entity_pair(self.sender_id)
self._sender, self._input_sender = utils._get_entity_pair(
self.sender_id, self._entities, client._entity_cache)
@property
def id(self):

View File

@ -206,16 +206,6 @@ class NewMessage(EventBuilder):
super()._set_client(client)
m = self.message
m._finish_init(client, self._entities, None)
# TODO Duplicated work here
m._chat, m._input_chat = self._get_entity_pair(m.chat_id)
m._sender, m._input_sender = self._get_entity_pair(m.sender_id)
m._via_bot, m._via_input_bot = self._get_entity_pair(m.via_bot_id)
if m.forward:
f = m.forward
f._chat, f._input_chat = self._get_entity_pair(f.chat_id)
f._sender, f._input_sender = self._get_entity_pair(f.sender_id)
self.__dict__['_init'] = True # No new attributes can be set
def __getattr__(self, item):

View File

@ -1,6 +1,7 @@
import datetime
from .common import EventBuilder, EventCommon, name_inner_event
from .. import utils
from ..tl import types
from ..tl.custom.sendergetter import SenderGetter
@ -184,7 +185,8 @@ class UserUpdate(EventBuilder):
self._chat_peer = types.PeerChannel(self._chat_peer)
super()._set_client(client)
self._sender, self._input_sender = self._get_entity_pair(self._sender_id)
self._sender, self._input_sender = utils._get_entity_pair(
self.sender_id, self._entities, client._entity_cache)
@property
def user(self):

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)

View File

@ -469,6 +469,24 @@ def get_input_message(message):
_raise_cast_fail(message, 'InputMedia')
def _get_entity_pair(entity_id, entities, cache,
get_input_peer=get_input_peer):
"""
Returns ``(entity, input_entity)`` for the given entity ID.
"""
entity = entities.get(entity_id)
try:
input_entity = cache[entity_id]
except KeyError:
# KeyError is unlikely, so another TypeError won't hurt
try:
input_entity = get_input_peer(entity)
except TypeError:
input_entity = None
return entity, input_entity
def get_message_id(message):
"""Similar to :meth:`get_input_peer`, but for message IDs."""
if message is None:

View File

@ -1,3 +1,3 @@
# Versions should comply with PEP440.
# This line is parsed in setup.py:
__version__ = '1.7.2'
__version__ = '1.7.3'