Add Spines class as a container for all Axes spines

timhoffm · y9c · timhoffm · commit c54a2e0f22bb · 2020-04-12T21:03:42.000+02:00
Co-authored-by: yech1990 &lt;yech1990@gmail.com&gt;
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -494,7 +494,7 @@ def __init__(self, fig, rect,
         self.set_box_aspect(box_aspect)
         self._axes_locator = None  # Optionally set via update(kwargs).
 
-        self.spines = self._gen_axes_spines()
+        self.spines = mspines.Spines.from_dict(self._gen_axes_spines())
 
         # this call may differ for non-sep axes, e.g., polar
         self._init_axis()
diff --git a/lib/matplotlib/spines.py b/lib/matplotlib/spines.py
@@ -1,3 +1,6 @@
+from collections.abc import MutableMapping
+import functools
+
 import numpy as np
 
 import matplotlib
@@ -534,3 +537,90 @@ def set_color(self, c):
         """
         self.set_edgecolor(c)
         self.stale = True
+
+
+class SpinesProxy:
+    """
+    A proxy to broadcast ``set_*`` method calls to all contained `.Spines`.
+
+    The supported methods are determined dynamically based on the contained
+    spines. If not all spines support a given method, it's executed only on
+    the subset of spines that support it.
+    """
+    def __init__(self, spine_dict):
+        self._spine_dict = spine_dict
+
+    def __getattr__(self, name):
+        broadcast_targets = [spine for spine in self._spine_dict.values()
+                             if hasattr(spine, name)]
+        if not name.startswith('set_') or not broadcast_targets:
+            raise ValueError(f"'SpinesProxy' object has no attribute '{name}'")
+
+        def x(_targets, _funcname, *args, **kwargs):
+            for spine in _targets:
+                getattr(spine, _funcname)(*args, **kwargs)
+        x = functools.partial(x, broadcast_targets, name)
+        x.__doc__ = broadcast_targets[0].__doc__
+        return x
+
+    def __dir__(self):
+        names = []
+        for spine in self._spine_dict.values():
+            names.extend(name
+                         for name in dir(spine) if name.startswith('set_'))
+        return list(sorted(set(names)))
+
+
+class Spines(MutableMapping):
+    r"""
+    The container of all `.Spine`\s in an Axes.
+
+    Individual spines can be accessed via attribute access of the spine names::
+
+        spines.top.set_visible(False)
+
+    For backward compatibility, it's also possible to use item access::
+
+        spines['top'].set_visible(False)
+
+    Setting spine properties can be broadcasted to all spines using the proxy
+    ``spines.all``:
+
+        spines.all.set_visible(False)
+
+    Attributes
+    ----------
+    all : SpinesProxy
+        A proxy to broadcast ``set_*`` method calls to all contained `.Spines`.
+
+    """
+    def __init__(self, **kwargs):
+        self._dict = kwargs
+        self.all = SpinesProxy(self._dict)
+
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**d)
+
+    def __getstate__(self):
+        return self._dict
+
+    def __setstate__(self, state):
+        self.__init__(**state)
+
+    def __getitem__(self, key):
+        return self._dict[key]
+
+    def __setitem__(self, key, value):
+        # TODO: Do we want to deprecate adding spines?
+        self._dict[key] = value
+
+    def __delitem__(self, key):
+        # TODO: Do we want to deprecate deleting spines?
+        del self._dict[key]
+
+    def __iter__(self):
+        return iter(self._dict)
+
+    def __len__(self):
+        return len(self._dict)
