Fix Bug in BasePersistence.insert/replace_bot for Objects with __dict__ in their slots (python-telegram-bot#2561)

Bibo-Joshi · web-flow · commit 52ce03929b32 · 2021-06-13T15:07:40.000+02:00
* Handle objects with __dict__ in __slots__

* Rework
diff --git a/telegram/ext/basepersistence.py b/telegram/ext/basepersistence.py
@@ -276,11 +276,8 @@ def _replace_bot(cls, obj: object, memo: Dict[int, object]) -> object:  # pylint
                 new_obj[cls._replace_bot(k, memo)] = cls._replace_bot(val, memo)
             memo[obj_id] = new_obj
             return new_obj
-        if hasattr(obj, '__dict__'):
-            for attr_name, attr in new_obj.__dict__.items():
-                setattr(new_obj, attr_name, cls._replace_bot(attr, memo))
-            memo[obj_id] = new_obj
-            return new_obj
+        # if '__dict__' in obj.__slots__, we already cover this here, that's why the
+        # __dict__ case comes below
         if hasattr(obj, '__slots__'):
             for attr_name in new_obj.__slots__:
                 setattr(
@@ -290,6 +287,11 @@ def _replace_bot(cls, obj: object, memo: Dict[int, object]) -> object:  # pylint
                 )
             memo[obj_id] = new_obj
             return new_obj
+        if hasattr(obj, '__dict__'):
+            for attr_name, attr in new_obj.__dict__.items():
+                setattr(new_obj, attr_name, cls._replace_bot(attr, memo))
+            memo[obj_id] = new_obj
+            return new_obj
 
         return obj
 
@@ -364,11 +366,8 @@ def _insert_bot(self, obj: object, memo: Dict[int, object]) -> object:  # pylint
                 new_obj[self._insert_bot(k, memo)] = self._insert_bot(val, memo)
             memo[obj_id] = new_obj
             return new_obj
-        if hasattr(obj, '__dict__'):
-            for attr_name, attr in new_obj.__dict__.items():
-                setattr(new_obj, attr_name, self._insert_bot(attr, memo))
-            memo[obj_id] = new_obj
-            return new_obj
+        # if '__dict__' in obj.__slots__, we already cover this here, that's why the
+        # __dict__ case comes below
         if hasattr(obj, '__slots__'):
             for attr_name in obj.__slots__:
                 setattr(
@@ -378,6 +377,11 @@ def _insert_bot(self, obj: object, memo: Dict[int, object]) -> object:  # pylint
                 )
             memo[obj_id] = new_obj
             return new_obj
+        if hasattr(obj, '__dict__'):
+            for attr_name, attr in new_obj.__dict__.items():
+                setattr(new_obj, attr_name, self._insert_bot(attr, memo))
+            memo[obj_id] = new_obj
+            return new_obj
 
         return obj
 
diff --git a/tests/test_persistence.py b/tests/test_persistence.py
@@ -561,14 +561,15 @@ class MyUpdate:
 
     def test_bot_replace_insert_bot(self, bot, bot_persistence):
         class CustomSlottedClass:
-            __slots__ = ('bot',)
+            __slots__ = ('bot', '__dict__')
 
             def __init__(self):
                 self.bot = bot
+                self.not_in_dict = bot
 
             def __eq__(self, other):
                 if isinstance(other, CustomSlottedClass):
-                    return self.bot is other.bot
+                    return self.bot is other.bot and self.not_in_dict is other.not_in_dict
                 return False
 
         class CustomClass:
@@ -587,6 +588,7 @@ def replace_bot():
                 cc = CustomClass()
                 cc.bot = BasePersistence.REPLACED_BOT
                 cc.slotted_object.bot = BasePersistence.REPLACED_BOT
+                cc.slotted_object.not_in_dict = BasePersistence.REPLACED_BOT
                 cc.list_ = [1, 2, BasePersistence.REPLACED_BOT]
                 cc.tuple_ = tuple(cc.list_)
                 cc.set_ = set(cc.list_)
