1717# You should have received a copy of the GNU Lesser Public License
1818# along with this program. If not, see [http://www.gnu.org/licenses/].
1919"""Base class for Telegram Objects."""
20+ from copy import deepcopy
21+
2022try :
2123 import ujson as json
2224except ImportError :
2325 import json # type: ignore[no-redef]
2426
25- from typing import TYPE_CHECKING , List , Optional , Type , TypeVar , Tuple
27+ from typing import TYPE_CHECKING , List , Optional , Type , TypeVar , Tuple , Dict , Union
2628
2729from telegram ._utils .types import JSONDict
2830from telegram ._utils .warnings import warn
@@ -40,6 +42,12 @@ class TelegramObject:
4042 is equivalent to ``telegram_object.attribute_name``. If the object does not have an attribute
4143 with the appropriate name, a :exc:`KeyError` will be raised.
4244
45+ When objects of this type are pickled, the :class:`~telegram.Bot` attribute associated with the
46+ object will be removed. However, when copying the object via :func:`copy.deepcopy`, the copy
47+ will have the *same* bot instance associated with it, i.e::
48+
49+ assert telegram_object.get_bot() is copy.deepcopy(telegram_object).get_bot()
50+
4351 .. versionchanged:: 14.0
4452 ``telegram_object['from']`` will look up the key ``from_user``. This is to account for
4553 special cases like :attr:`Message.from_user` that deviate from the official Bot API.
@@ -53,15 +61,12 @@ class TelegramObject:
5361 _bot : Optional ['Bot' ]
5462 # Adding slots reduces memory usage & allows for faster attribute access.
5563 # Only instance variables should be added to __slots__.
56- __slots__ = (
57- '_id_attrs' ,
58- '_bot' ,
59- )
64+ __slots__ = ('_id_attrs' , '_bot' )
6065
6166 # pylint: disable=unused-argument
6267 def __new__ (cls , * args : object , ** kwargs : object ) -> 'TelegramObject' :
6368 # We add _id_attrs in __new__ instead of __init__ since we want to add this to the slots
64- # w/o calling __init__ in all of the subclasses. This is what we also do in BaseFilter.
69+ # w/o calling __init__ in all of the subclasses.
6570 instance = super ().__new__ (cls )
6671 instance ._id_attrs = ()
6772 instance ._bot = None
@@ -81,6 +86,86 @@ def __getitem__(self, item: str) -> object:
8186 f"`{ item }
8287 ) from exc
8388
89+ def __getstate__ (self ) -> Dict [str , Union [str , object ]]:
90+ """
91+ This method is used for pickling. We remove the bot attribute of the object since those
92+ are not pickable.
93+ """
94+ return self ._get_attrs (include_private = True , recursive = False , remove_bot = True )
95+
96+ def __setstate__ (self , state : dict ) -> None :
97+ """
98+ This method is used for unpickling. The data, which is in the form a dictionary, is
99+ converted back into a class. Should be modified in place.
100+ """
101+ for key , val in state .items ():
102+ setattr (self , key , val )
103+
104+ def __deepcopy__ (self : TO , memodict : dict ) -> TO :
105+ """This method deepcopies the object and sets the bot on the newly created copy."""
106+ bot = self ._bot # Save bot so we can set it after copying
107+ self .set_bot (None ) # set to None so it is not deepcopied
108+ cls = self .__class__
109+ result = cls .__new__ (cls ) # create a new instance
110+ memodict [id (self )] = result # save the id of the object in the dict
111+
112+ attrs = self ._get_attrs (include_private = True ) # get all its attributes
113+
114+ for k in attrs : # now we set the attributes in the deepcopied object
115+ setattr (result , k , deepcopy (getattr (self , k ), memodict ))
116+
117+ result .set_bot (bot ) # Assign the bots back
118+ self .set_bot (bot )
119+ return result # type: ignore[return-value]
120+
121+ def _get_attrs (
122+ self ,
123+ include_private : bool = False ,
124+ recursive : bool = False ,
125+ remove_bot : bool = False ,
126+ ) -> Dict [str , Union [str , object ]]:
127+ """This method is used for obtaining the attributes of the object.
128+
129+ Args:
130+ include_private (:obj:`bool`): Whether the result should include private variables.
131+ recursive (:obj:`bool`): If :obj:`True`, will convert any TelegramObjects (if found) in
132+ the attributes to a dictionary. Else, preserves it as an object itself.
133+ remove_bot (:obj:`bool`): Whether the bot should be included in the result.
134+
135+ Returns:
136+ :obj:`dict`: A dict where the keys are attribute names and values are their values.
137+ """
138+ data = {}
139+
140+ if not recursive :
141+ try :
142+ # __dict__ has attrs from superclasses, so no need to put in the for loop below
143+ data .update (self .__dict__ )
144+ except AttributeError :
145+ pass
146+ # We want to get all attributes for the class, using self.__slots__ only includes the
147+ # attributes used by that class itself, and not its superclass(es). Hence, we get its MRO
148+ # and then get their attributes. The `[:-1]` slice excludes the `object` class
149+ for cls in self .__class__ .__mro__ [:- 1 ]:
150+ for key in cls .__slots__ :
151+ if not include_private and key .startswith ('_' ):
152+ continue
153+
154+ value = getattr (self , key , None )
155+ if value is not None :
156+ if recursive and hasattr (value , 'to_dict' ):
157+ data [key ] = value .to_dict ()
158+ else :
159+ data [key ] = value
160+ elif not recursive :
161+ data [key ] = value
162+
163+ if recursive and data .get ('from_user' ):
164+ data ['from' ] = data .pop ('from_user' , None )
165+ if remove_bot :
166+ data .pop ('_bot' , None )
167+ return data
168+
84169 @staticmethod
85170 def _parse_data (data : Optional [JSONDict ]) -> Optional [JSONDict ]:
86171 return None if data is None else data .copy ()
@@ -137,27 +222,7 @@ def to_dict(self) -> JSONDict:
137222 Returns:
138223 :obj:`dict`
139224 """
140- data = {}
141-
142- # We want to get all attributes for the class, using self.__slots__ only includes the
143- # attributes used by that class itself, and not its superclass(es). Hence we get its MRO
144- # and then get their attributes. The `[:-2]` slice excludes the `object` class & the
145- # TelegramObject class itself.
146- attrs = {attr for cls in self .__class__ .__mro__ [:- 2 ] for attr in cls .__slots__ }
147- for key in attrs :
148- if key == 'bot' or key .startswith ('_' ):
149- continue
150-
151- value = getattr (self , key , None )
152- if value is not None :
153- if hasattr (value , 'to_dict' ):
154- data [key ] = value .to_dict ()
155- else :
156- data [key ] = value
157-
158- if data .get ('from_user' ):
159- data ['from' ] = data .pop ('from_user' , None )
160- return data
225+ return self ._get_attrs (recursive = True )
161226
162227 def get_bot (self ) -> 'Bot' :
163228 """Returns the :class:`telegram.Bot` instance associated with this object.
@@ -171,8 +236,7 @@ def get_bot(self) -> 'Bot':
171236 """
172237 if self ._bot is None :
173238 raise RuntimeError (
174- 'This object has no bot associated with it. \
175-
239+ 'This object has no bot associated with it. Shortcuts cannot be used.'
176240 )
177241 return self ._bot
178242
0 commit comments