Improve Backwards Compatibility of TelegramObjects Pickle Behavior (#3382)

Co-authored-by: Hinrich Mahler <22366557+Bibo-Joshi@users.noreply.github.com>
This commit is contained in:
Harshil 2022-11-24 16:41:37 +05:30 committed by GitHub
parent 05c6ca06f8
commit 1724212458
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 101 additions and 16 deletions

View file

@ -22,7 +22,19 @@ import inspect
import json import json
from copy import deepcopy from copy import deepcopy
from itertools import chain 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.datetime import to_timestamp
from telegram._utils.types import JSONDict from telegram._utils.types import JSONDict
@ -183,12 +195,16 @@ class TelegramObject:
""" """
Overrides :meth:`object.__setstate__` to customize the unpickling process of objects of Overrides :meth:`object.__setstate__` to customize the unpickling process of objects of
this type. Modifies the object in-place. 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 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 if the class now has dedicated attributes for those keys and moves the values from
:attr:`api_kwargs` to the dedicated attributes. :attr:`api_kwargs` to the dedicated attributes.
This can happen, if serialized data is loaded with a new version of this library, where 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. 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: Args:
state (:obj:`dict`): The data to set as attributes of this object. state (:obj:`dict`): The data to set as attributes of this object.
""" """
@ -196,8 +212,14 @@ class TelegramObject:
# this as Bots are not pickable. # this as Bots are not pickable.
setattr(self, "_bot", None) setattr(self, "_bot", None)
setattr(self, "api_kwargs", state.pop("api_kwargs", {})) # assign api_kwargs first
for key, val in state.items(): 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() self._apply_api_kwargs()
def __deepcopy__(self: Tele_co, memodict: dict) -> Tele_co: def __deepcopy__(self: Tele_co, memodict: dict) -> Tele_co:
@ -221,15 +243,45 @@ class TelegramObject:
result = cls.__new__(cls) # create a new instance result = cls.__new__(cls) # create a new instance
memodict[id(self)] = result # save the id of the object in the dict 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 self._get_attrs_names(
include_private=True
for k in attrs: # now we set the attributes in the deepcopied object ): # now we set the attributes in the deepcopied object
setattr(result, k, deepcopy(getattr(self, k), memodict)) 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 result.set_bot(bot) # Assign the bots back
self.set_bot(bot) self.set_bot(bot)
return result 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( def _get_attrs(
self, self,
include_private: bool = False, include_private: bool = False,
@ -249,14 +301,7 @@ class TelegramObject:
""" """
data = {} data = {}
# We want to get all attributes for the class, using self.__slots__ only includes the for key in self._get_attrs_names(include_private=include_private):
# 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
value = getattr(self, key, None) value = getattr(self, key, None)
if value is not None: if value is not None:

View file

@ -248,7 +248,7 @@ PROJECT_ROOT_PATH = Path(__file__).parent.parent.resolve()
TEST_DATA_PATH = Path(__file__).parent.resolve() / "data" TEST_DATA_PATH = Path(__file__).parent.resolve() / "data"
def data_file(filename: str): def data_file(filename: str) -> Path:
return TEST_DATA_PATH / filename return TEST_DATA_PATH / filename

Binary file not shown.

View file

@ -25,6 +25,8 @@ from pathlib import Path
import pytest import pytest
from telegram import Bot, BotCommand, Chat, Message, PhotoSize, TelegramObject, User 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): def all_subclasses(cls):
@ -138,6 +140,18 @@ class TestTelegramObject:
to = TelegramObject(api_kwargs={"foo": "bar"}) to = TelegramObject(api_kwargs={"foo": "bar"})
assert to.to_dict() == {"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): def test_to_dict_recursion(self):
class Recursive(TelegramObject): class Recursive(TelegramObject):
__slots__ = ("recursive",) __slots__ = ("recursive",)
@ -243,7 +257,7 @@ class TestTelegramObject:
assert unpickled.date == date, f"{unpickled.date} != {date}" assert unpickled.date == date, f"{unpickled.date} != {date}"
assert unpickled.photo[0] == photo 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 """Makes sure that when a class gets new attributes, the api_kwargs are moved to the
new attributes on unpickling.""" new attributes on unpickling."""
obj = self.ChangingTO(api_kwargs={"foo": "bar"}) obj = self.ChangingTO(api_kwargs={"foo": "bar"})
@ -255,6 +269,32 @@ class TestTelegramObject:
assert obj.foo == "bar" assert obj.foo == "bar"
assert obj.api_kwargs == {} 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): def test_deepcopy_telegram_obj(self, bot):
chat = Chat(2, Chat.PRIVATE) chat = Chat(2, Chat.PRIVATE)
user = User(3, "first_name", False) user = User(3, "first_name", False)