Handle Classes as Input for BasePersistence.replace/insert_bot (python-telegram-bot#2523)

Bibo-Joshi · web-flow · commit 08ba7c7793dc · 2021-05-16T20:02:35.000+02:00
* Ignore classes on replace/insert_bot

* Review
diff --git a/telegram/ext/basepersistence.py b/telegram/ext/basepersistence.py
@@ -137,7 +137,8 @@ def replace_bot(cls, obj: object) -> object:
         Replaces all instances of :class:`telegram.Bot` that occur within the passed object with
         :attr:`REPLACED_BOT`. Currently, this handles objects of type ``list``, ``tuple``, ``set``,
         ``frozenset``, ``dict``, ``defaultdict`` and objects that have a ``__dict__`` or
-        ``__slot__`` attribute, excluding objects that can't be copied with `copy.copy`.
+        ``__slot__`` attribute, excluding classes and objects that can't be copied with
+        ``copy.copy``.
 
         Args:
             obj (:obj:`object`): The object
@@ -168,6 +169,14 @@ def _replace_bot(cls, obj: object, memo: Dict[int, object]) -> object:  # pylint
             new_immutable = obj.__class__(cls._replace_bot(item, memo) for item in obj)
             memo[obj_id] = new_immutable
             return new_immutable
+        if isinstance(obj, type):
+            # classes usually do have a __dict__, but it's not writable
+            warnings.warn(
+                'BasePersistence.replace_bot does not handle classes. See '
+                'the docs of BasePersistence.replace_bot for more information.',
+                RuntimeWarning,
+            )
+            return obj
 
         try:
             new_obj = copy(obj)
@@ -215,7 +224,8 @@ def insert_bot(self, obj: object) -> object:
         Replaces all instances of :attr:`REPLACED_BOT` that occur within the passed object with
         :attr:`bot`. Currently, this handles objects of type ``list``, ``tuple``, ``set``,
         ``frozenset``, ``dict``, ``defaultdict`` and objects that have a ``__dict__`` or
-        ``__slot__`` attribute, excluding objects that can't be copied with `copy.copy`.
+        ``__slot__`` attribute, excluding classes and objects that can't be copied with
+        ``copy.copy``.
 
         Args:
             obj (:obj:`object`): The object
@@ -248,6 +258,14 @@ def _insert_bot(self, obj: object, memo: Dict[int, object]) -> object:  # pylint
             new_immutable = obj.__class__(self._insert_bot(item, memo) for item in obj)
             memo[obj_id] = new_immutable
             return new_immutable
+        if isinstance(obj, type):
+            # classes usually do have a __dict__, but it's not writable
+            warnings.warn(
+                'BasePersistence.insert_bot does not handle classes. See '
+                'the docs of BasePersistence.insert_bot for more information.',
+                RuntimeWarning,
+            )
+            return obj
 
         try:
             new_obj = copy(obj)
diff --git a/tests/test_persistence.py b/tests/test_persistence.py
@@ -590,6 +590,33 @@ def __copy__(self):
             "BasePersistence.insert_bot does not handle objects that can not be copied."
         )
 
+    def test_bot_replace_insert_bot_classes(self, bot, bot_persistence, recwarn):
+        """Here check that classes are just returned verbatim."""
+        persistence = bot_persistence
+        persistence.set_bot(bot)
+
+        class CustomClass:
+            pass
+
+        persistence.update_bot_data({1: CustomClass})
+        assert persistence.bot_data[1] is CustomClass
+        persistence.update_chat_data(123, {1: CustomClass})
+        assert persistence.chat_data[123][1] is CustomClass
+        persistence.update_user_data(123, {1: CustomClass})
+        assert persistence.user_data[123][1] is CustomClass
+
+        assert persistence.get_bot_data()[1] is CustomClass
+        assert persistence.get_chat_data()[123][1] is CustomClass
+        assert persistence.get_user_data()[123][1] is CustomClass
+
+        assert len(recwarn) == 2
+        assert str(recwarn[0].message).startswith(
+            "BasePersistence.replace_bot does not handle classes."
+        )
+        assert str(recwarn[1].message).startswith(
+            "BasePersistence.insert_bot does not handle classes."
+        )
+
     def test_bot_replace_insert_bot_objects_with_faulty_equality(self, bot, bot_persistence):
         """Here check that trying to compare obj == self.REPLACED_BOT doesn't lead to problems."""
         persistence = bot_persistence
