Improve Backwards Compatibility of TelegramObjects Pickle Behavior (#3382)

Co-authored-by: Hinrich Mahler <22366557+Bibo-Joshi@users.noreply.github.com>

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

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
 
 

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)