Improve Subscription of TelegramObject (#2719)

Co-authored-by: Hinrich Mahler <22366557+Bibo-Joshi@users.noreply.github.com>
Co-authored-by: Harshil <37377066+harshil21@users.noreply.github.com>
This commit is contained in:
Simon Damberg 2021-11-21 13:36:22 +01:00 committed by Hinrich Mahler
parent 55e8c38e50
commit 28c8c155ce
4 changed files with 35 additions and 12 deletions

View file

@ -672,16 +672,13 @@ class Message(TelegramObject):
for attachment_type in MessageAttachmentType:
if self[attachment_type]:
self._effective_attachment = self[attachment_type]
self._effective_attachment = self[attachment_type] # type: ignore[assignment]
break
else:
self._effective_attachment = None
return self._effective_attachment # type: ignore[return-value]
def __getitem__(self, item: str) -> Any: # pylint: disable=inconsistent-return-statements
return self.chat.id if item == 'chat_id' else super().__getitem__(item)
def to_dict(self) -> JSONDict:
"""See :meth:`telegram.TelegramObject.to_dict`."""
data = super().to_dict()

View file

@ -34,7 +34,16 @@ TO = TypeVar('TO', bound='TelegramObject', covariant=True)
class TelegramObject:
"""Base class for most Telegram objects."""
"""Base class for most Telegram objects.
Objects of this type are subscriptable with strings, where ``telegram_object[attribute_name]``
is equivalent to ``telegram_object.attribute_name``. If the object does not have an attribute
with the appropriate name, a :exc:`KeyError` will be raised.
.. versionchanged:: 14.0
``telegram_object['from']`` will look up the key ``from_user``. This is to account for
special cases like :attr:`Message.from_user` that deviate from the official Bot API.
"""
# type hints in __new__ are not read by mypy (https://github.com/python/mypy/issues/1021). As a
# workaround we can type hint instance variables in __new__ using a syntax defined in PEP 526 -
@ -62,7 +71,15 @@ class TelegramObject:
return str(self.to_dict())
def __getitem__(self, item: str) -> object:
return getattr(self, item, None)
if item == 'from':
item = 'from_user'
try:
return getattr(self, item)
except AttributeError as exc:
raise KeyError(
f"Objects of type {self.__class__.__name__} don't have an attribute called "
f"`{item}`."
) from exc
@staticmethod
def _parse_data(data: Optional[JSONDict]) -> Optional[JSONDict]:

View file

@ -321,11 +321,6 @@ class TestMessage:
assert new.to_dict() == message_params.to_dict()
def test_dict_approach(self, message):
assert message['text'] == message.text
assert message['chat_id'] == message.chat_id
assert message['no_key'] is None
def test_parse_entity(self):
text = (
b'\\U0001f469\\u200d\\U0001f469\\u200d\\U0001f467'

View file

@ -26,7 +26,7 @@ try:
except ImportError:
ujson = None
from telegram import TelegramObject
from telegram import TelegramObject, Message, Chat, User
class TestTelegramObject:
@ -131,3 +131,17 @@ class TestTelegramObject:
elif bot_inst is None:
with pytest.raises(RuntimeError):
tg_object.get_bot()
def test_subscription(self):
# We test with Message because that gives us everything we want to test - easier than
# implementing a custom subclass just for this test
chat = Chat(2, Chat.PRIVATE)
user = User(3, 'first_name', False)
message = Message(1, None, chat=chat, from_user=user, text='foobar')
assert message['text'] == 'foobar'
assert message['chat'] is chat
assert message['chat_id'] == 2
assert message['from'] is user
assert message['from_user'] is user
with pytest.raises(KeyError, match="Message don't have an attribute called `no_key`"):
message['no_key']