Merge pull request #17620 from tacaswell/simplified_setattr_cm_method_handling

anntzer · web-flow · commit e86e05960610 · 2020-06-15T19:26:21.000+02:00
MNT: make _setattr_cm more conservative
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -2032,13 +2032,30 @@ def _setattr_cm(obj, **kwargs):
     Temporarily set some attributes; restore original state at context exit.
     """
     sentinel = object()
-    origs = [(attr, getattr(obj, attr, sentinel)) for attr in kwargs]
+    origs = {}
+    for attr in kwargs:
+        orig = getattr(obj, attr, sentinel)
+
+        if attr in obj.__dict__ or orig is sentinel:
+            origs[attr] = orig
+        else:
+            cls_orig = getattr(type(obj), attr)
+            if isinstance(cls_orig, property):
+                origs[attr] = orig
+            elif isinstance(cls_orig, types.FunctionType):
+                origs[attr] = sentinel
+            else:
+                raise ValueError(
+                    f"trying to set {attr} which is not a method, "
+                    "property, or instance level attribute"
+                )
+
     try:
         for attr, val in kwargs.items():
             setattr(obj, attr, val)
         yield
     finally:
-        for attr, orig in origs:
+        for attr, orig in origs.items():
             if orig is sentinel:
                 delattr(obj, attr)
             else:
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -651,3 +651,73 @@ def test_check_shape(target, test_shape):
     with pytest.raises(ValueError,
                        match=error_pattern):
         cbook._check_shape(target, aardvark=data)
+
+
+def test_setattr_cm():
+    class A:
+
+        cls_level = object()
+        override = object()
+        def __init__(self):
+            self.aardvark = 'aardvark'
+            self.override = 'override'
+            self._p = 'p'
+
+        def meth(self):
+            ...
+
+        @property
+        def prop(self):
+            return self._p
+
+        @prop.setter
+        def prop(self, val):
+            self._p = val
+
+    class B(A):
+        ...
+
+    other = A()
+
+    def verify_pre_post_state(obj):
+        # When you access a Python method the function is bound
+        # to the object at access time so you get a new instance
+        # of MethodType every time.
+        #
+        # https://docs.python.org/3/howto/descriptor.html#functions-and-methods
+        assert obj.meth is not obj.meth
+        # normal attribute should give you back the same instance every time
+        assert obj.aardvark is obj.aardvark
+        assert a.aardvark == 'aardvark'
+        # and our property happens to give the same instance every time
+        assert obj.prop is obj.prop
+        assert obj.cls_level is A.cls_level
+        assert obj.override == 'override'
+        assert not hasattr(obj, 'extra')
+        assert obj.prop == 'p'
+        assert obj.monkey == other.meth
+
+    a = B()
+
+    a.monkey = other.meth
+    verify_pre_post_state(a)
+    with cbook._setattr_cm(
+            a, prop='squirrel',
+            aardvark='moose', meth=lambda: None,
+            override='boo', extra='extra',
+            monkey=lambda: None):
+        # because we have set a lambda, it is normal attribute access
+        # and the same every time
+        assert a.meth is a.meth
+        assert a.aardvark is a.aardvark
+        assert a.aardvark == 'moose'
+        assert a.override == 'boo'
+        assert a.extra == 'extra'
+        assert a.prop == 'squirrel'
+        assert a.monkey != other.meth
+
+    verify_pre_post_state(a)
+
+    with pytest.raises(ValueError):
+        with cbook._setattr_cm(a, cls_level='bob'):
+            pass
