Handle Properties in TelegramObject.__setstate__ (#4134)

Co-authored-by: Bibo-Joshi <22366557+Bibo-Joshi@users.noreply.github.com>
This commit is contained in:
Harshil 2024-03-03 13:22:26 -05:00 committed by GitHub
parent 5d11d7fd42
commit bd9b0bd126
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 83 additions and 2 deletions

View file

@ -17,6 +17,7 @@
# You should have received a copy of the GNU Lesser Public License # You should have received a copy of the GNU Lesser Public License
# along with this program. If not, see [http://www.gnu.org/licenses/]. # along with this program. If not, see [http://www.gnu.org/licenses/].
"""Base class for Telegram Objects.""" """Base class for Telegram Objects."""
import contextlib
import datetime import datetime
import inspect import inspect
import json import json
@ -312,7 +313,20 @@ class TelegramObject:
try: try:
setattr(self, key, val) setattr(self, key, val)
except AttributeError: except AttributeError:
# catch cases when old attributes are removed from new versions # So an attribute was deprecated and removed from the class. Let's handle this:
# 1) Is the attribute now a property with no setter? Let's check that:
if isinstance(getattr(self.__class__, key, None), property):
# It is, so let's try to set the "private attribute" instead
try:
setattr(self, f"_{key}", val)
# If this fails as well, guess we've completely removed it. Let's add it to
# api_kwargs as fallback
except AttributeError:
api_kwargs[key] = val
# 2) The attribute is a private attribute, i.e. it went through case 1) in the past
elif key.startswith("_"):
continue # skip adding this to api_kwargs, the attribute is lost forever.
api_kwargs[key] = val # add it to api_kwargs as fallback api_kwargs[key] = val # add it to api_kwargs as fallback
# For api_kwargs we first apply any kwargs that are already attributes of the object # For api_kwargs we first apply any kwargs that are already attributes of the object
@ -490,7 +504,12 @@ class TelegramObject:
""" """
# we convert to list to ensure that the list doesn't change length while we loop # we convert to list to ensure that the list doesn't change length while we loop
for key in list(api_kwargs.keys()): for key in list(api_kwargs.keys()):
if getattr(self, key, True) is None: # property attributes are not settable, so we need to set the private attribute
if isinstance(getattr(self.__class__, key, None), property):
# if setattr fails, we'll just leave the value in api_kwargs:
with contextlib.suppress(AttributeError):
setattr(self, f"_{key}", api_kwargs.pop(key))
elif getattr(self, key, True) is None:
setattr(self, key, api_kwargs.pop(key)) setattr(self, key, api_kwargs.pop(key))
def _get_attrs_names(self, include_private: bool) -> Iterator[str]: def _get_attrs_names(self, include_private: bool) -> Iterator[str]:

View file

@ -358,6 +358,68 @@ class TestTelegramObject:
chat.id = 7 chat.id = 7
assert chat.id == 7 assert chat.id == 7
def test_pickle_handle_properties(self):
# Very hard to properly test, can't use a pickle file since newer versions of the library
# will stop having the property.
# The code below uses exec statements to simulate library changes. There is no other way
# to test this.
# Original class:
v1 = """
class PicklePropertyTest(TelegramObject):
__slots__ = ("forward_from", "to_be_removed", "forward_date")
def __init__(self, forward_from=None, forward_date=None, api_kwargs=None):
super().__init__(api_kwargs=api_kwargs)
self.forward_from = forward_from
self.forward_date = forward_date
self.to_be_removed = "to_be_removed"
"""
exec(v1, globals(), None)
old = PicklePropertyTest("old_val", "date", api_kwargs={"new_attr": 1}) # noqa: F821
pickled_v1 = pickle.dumps(old)
# After some API changes:
v2 = """
class PicklePropertyTest(TelegramObject):
__slots__ = ("_forward_from", "_date", "_new_attr")
def __init__(self, forward_from=None, f_date=None, new_attr=None, api_kwargs=None):
super().__init__(api_kwargs=api_kwargs)
self._forward_from = forward_from
self.f_date = f_date
self._new_attr = new_attr
@property
def forward_from(self):
return self._forward_from
@property
def forward_date(self):
return self.f_date
@property
def new_attr(self):
return self._new_attr
"""
exec(v2, globals(), None)
v2_unpickle = pickle.loads(pickled_v1)
assert v2_unpickle.forward_from == "old_val" == v2_unpickle._forward_from
with pytest.raises(AttributeError):
# New attribute should not be available either as is always the case for pickle
v2_unpickle.forward_date
assert v2_unpickle.new_attr == 1 == v2_unpickle._new_attr
assert not hasattr(v2_unpickle, "to_be_removed")
assert v2_unpickle.api_kwargs == {"to_be_removed": "to_be_removed"}
pickled_v2 = pickle.dumps(v2_unpickle)
# After PTB removes the property and the attribute:
v3 = """
class PicklePropertyTest(TelegramObject):
__slots__ = ()
def __init__(self, api_kwargs=None):
super().__init__(api_kwargs=api_kwargs)
"""
exec(v3, globals(), None)
v3_unpickle = pickle.loads(pickled_v2)
assert v3_unpickle.api_kwargs == {"to_be_removed": "to_be_removed"}
assert not hasattr(v3_unpickle, "_forward_from")
assert not hasattr(v3_unpickle, "_new_attr")
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)