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):