diff --git a/telegram/_telegramobject.py b/telegram/_telegramobject.py index a9a31cd3c..6ca8c8af3 100644 --- a/telegram/_telegramobject.py +++ b/telegram/_telegramobject.py @@ -22,7 +22,19 @@ import inspect import json from copy import deepcopy from itertools import chain -from typing import TYPE_CHECKING, Dict, List, Optional, Set, Sized, Tuple, Type, TypeVar, Union +from typing import ( + TYPE_CHECKING, + Dict, + Iterator, + List, + Optional, + Set, + Sized, + Tuple, + Type, + TypeVar, + Union, +) from telegram._utils.datetime import to_timestamp from telegram._utils.types import JSONDict @@ -183,12 +195,16 @@ class TelegramObject: """ Overrides :meth:`object.__setstate__` to customize the unpickling process of objects of this type. Modifies the object in-place. + If any data was stored in the :attr:`api_kwargs` of the pickled object, this method checks if the class now has dedicated attributes for those keys and moves the values from :attr:`api_kwargs` to the dedicated attributes. This can happen, if serialized data is loaded with a new version of this library, where the new version was updated to account for updates of the Telegram Bot API. + If on the contrary an attribute was removed from the class, the value is not discarded but + made available via :attr:`api_kwargs`. + Args: state (:obj:`dict`): The data to set as attributes of this object. """ @@ -196,8 +212,14 @@ class TelegramObject: # this as Bots are not pickable. setattr(self, "_bot", None) + setattr(self, "api_kwargs", state.pop("api_kwargs", {})) # assign api_kwargs first + for key, val in state.items(): - setattr(self, key, val) + try: + setattr(self, key, val) + except AttributeError: # catch cases when old attributes are removed from new versions + self.api_kwargs[key] = val # add it to api_kwargs as fallback + self._apply_api_kwargs() def __deepcopy__(self: Tele_co, memodict: dict) -> Tele_co: @@ -221,15 +243,45 @@ class TelegramObject: result = cls.__new__(cls) # create a new instance memodict[id(self)] = result # save the id of the object in the dict - attrs = self._get_attrs(include_private=True) # get all its attributes - - for k in attrs: # now we set the attributes in the deepcopied object - setattr(result, k, deepcopy(getattr(self, k), memodict)) + for k in self._get_attrs_names( + include_private=True + ): # now we set the attributes in the deepcopied object + try: + setattr(result, k, deepcopy(getattr(self, k), memodict)) + except AttributeError: + # Skip missing attributes. This can happen if the object was loaded from a pickle + # file that was created with an older version of the library, where the class + # did not have the attribute yet. + continue result.set_bot(bot) # Assign the bots back self.set_bot(bot) return result + def _get_attrs_names(self, include_private: bool) -> Iterator[str]: + """ + Returns the names of the attributes of this object. This is used to determine which + attributes should be serialized when pickling the object. + + Args: + include_private (:obj:`bool`): Whether to include private attributes. + + Returns: + Iterator[:obj:`str`]: An iterator over the names of the attributes of this object. + """ + # We want to get all attributes for the class, using self.__slots__ only includes the + # attributes used by that class itself, and not its superclass(es). Hence, we get its MRO + # and then get their attributes. The `[:-1]` slice excludes the `object` class + all_slots = (s for c in self.__class__.__mro__[:-1] for s in c.__slots__) # type: ignore + # chain the class's slots with the user defined subclass __dict__ (class has no slots) + all_attrs = ( + chain(all_slots, self.__dict__.keys()) if hasattr(self, "__dict__") else all_slots + ) + + if include_private: + return all_attrs + return (attr for attr in all_attrs if not attr.startswith("_")) + def _get_attrs( self, include_private: bool = False, @@ -249,14 +301,7 @@ class TelegramObject: """ data = {} - # We want to get all attributes for the class, using self.__slots__ only includes the - # attributes used by that class itself, and not its superclass(es). Hence, we get its MRO - # and then get their attributes. The `[:-1]` slice excludes the `object` class - all_slots = (s for c in self.__class__.__mro__[:-1] for s in c.__slots__) # type: ignore - # chain the class's slots with the user defined subclass __dict__ (class has no slots) - for key in chain(self.__dict__, all_slots) if hasattr(self, "__dict__") else all_slots: - if not include_private and key.startswith("_"): - continue + for key in self._get_attrs_names(include_private=include_private): value = getattr(self, key, None) if value is not None: diff --git a/tests/conftest.py b/tests/conftest.py index 9d56aab76..e05442dad 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -248,7 +248,7 @@ PROJECT_ROOT_PATH = Path(__file__).parent.parent.resolve() TEST_DATA_PATH = Path(__file__).parent.resolve() / "data" -def data_file(filename: str): +def data_file(filename: str) -> Path: return TEST_DATA_PATH / filename diff --git a/tests/data/20a5_modified_chat.pickle b/tests/data/20a5_modified_chat.pickle new file mode 100644 index 000000000..49e75ac60 Binary files /dev/null and b/tests/data/20a5_modified_chat.pickle differ diff --git a/tests/test_telegramobject.py b/tests/test_telegramobject.py index 5eb248200..02549b06c 100644 --- a/tests/test_telegramobject.py +++ b/tests/test_telegramobject.py @@ -25,6 +25,8 @@ from pathlib import Path import pytest from telegram import Bot, BotCommand, Chat, Message, PhotoSize, TelegramObject, User +from telegram.ext import PicklePersistence +from tests.conftest import data_file def all_subclasses(cls): @@ -138,6 +140,18 @@ class TestTelegramObject: to = TelegramObject(api_kwargs={"foo": "bar"}) assert to.to_dict() == {"foo": "bar"} + def test_to_dict_missing_attribute(self): + message = Message( + 1, datetime.datetime.now(), Chat(1, "private"), from_user=User(1, "", False) + ) + del message.chat + + message_dict = message.to_dict() + assert "chat" not in message_dict + + message_dict = message.to_dict(recursive=False) + assert message_dict["chat"] is None + def test_to_dict_recursion(self): class Recursive(TelegramObject): __slots__ = ("recursive",) @@ -243,7 +257,7 @@ class TestTelegramObject: assert unpickled.date == date, f"{unpickled.date} != {date}" assert unpickled.photo[0] == photo - def test_pickle_apply_api_kwargs(self, bot): + def test_pickle_apply_api_kwargs(self): """Makes sure that when a class gets new attributes, the api_kwargs are moved to the new attributes on unpickling.""" obj = self.ChangingTO(api_kwargs={"foo": "bar"}) @@ -255,6 +269,32 @@ class TestTelegramObject: assert obj.foo == "bar" assert obj.api_kwargs == {} + async def test_pickle_removed_and_added_attribute(self): + """Test when newer versions of the library remove or add attributes from classes (which + the old pickled versions still/don't have). + """ + # We use a modified version of the 20.0a5 Chat class, which + # * has an `all_members_are_admins` attribute, + # * a non-empty `api_kwargs` dict + # * does not have the `is_forum` attribute + # This specific version was pickled + # using PicklePersistence.update_chat_data and that's what we use here to test if + # * the (now) removed attribute `all_members_are_admins` was added to api_kwargs + # * the (now) added attribute `is_forum` does not affect the unpickling + pp = PicklePersistence(data_file("20a5_modified_chat.pickle")) + chat = (await pp.get_chat_data())[1] + assert chat.id == 1 and chat.type == Chat.PRIVATE + assert chat.api_kwargs == { + "all_members_are_administrators": True, + "something": "Manually inserted", + } + with pytest.raises(AttributeError): + # removed attribute should not be available as attribute, only though api_kwargs + chat.all_members_are_administrators + with pytest.raises(AttributeError): + # New attribute should not be available either as is always the case for pickle + chat.is_forum + def test_deepcopy_telegram_obj(self, bot): chat = Chat(2, Chat.PRIVATE) user = User(3, "first_name", False)