MNT: make _setattr_cm more conservative

tacaswell · tacaswell · commit 4111ab61daf7 · 2020-06-11T17:11:32.000-04:00
- special case methods because they are descriptors
 - fail on any other non-instance attribute
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -2031,14 +2031,34 @@ def _setattr_cm(obj, **kwargs):
     """
     Temporarily set some attributes; restore original state at context exit.
     """
+    import inspect
     sentinel = object()
-    origs = [(attr, getattr(obj, attr, sentinel)) for attr in kwargs]
+    origs = {}
+    for attr in kwargs:
+        orig = getattr(obj, attr, sentinel)
+        # monkey patching on a new attribute, this is OK
+        if orig is sentinel:
+            origs[attr] = sentinel
+        # special case methods
+        elif inspect.ismethod(orig):
+            origs[attr] = sentinel
+        else:
+            # if we are trying to monkey patch a non-instance attribute,
+            # fail so we don't have to sort out how to safely identify
+            # general descriptors.
+            if attr not in obj.__dict__:
+                raise ValueError(
+                    f"trying to set {attr} which is not a method or "
+                    "instance level attribute"
+                )
+            origs[attr] = orig
+
     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,69 @@ 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):
+        ...
+
+    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')
+
+    a = B()
+    verify_pre_post_state(a)
+    with cbook._setattr_cm(
+            a,
+            aardvark='moose', meth=lambda: None,
+            override='boo', extra='extra'
+    ):
+        # 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'
+    verify_pre_post_state(a)
+
+    with pytest.raises(ValueError):
+        with cbook._setattr_cm(a, prop=''):
+            ...
+
+    with pytest.raises(ValueError):
+        with cbook._setattr_cm(a, cls_level='bob'):
+            ...
