Add Spines class as a container for all Axes spines

timhoffm · y9c · timhoffm · commit 6229785a1a05 · 2020-10-20T17:00:19.000+02:00
Co-authored-by: yech1990 &lt;yech1990@gmail.com&gt;
diff --git a/doc/users/next_whats_new/2020-04-12-spines.rst b/doc/users/next_whats_new/2020-04-12-spines.rst
@@ -0,0 +1,19 @@
+``Axes.spines``
+---------------
+
+``Axes.spines`` is now a dedicated container class `.Spines` for a set of
+`.Spine`\s instead of an ``OrderedDict``. On top of dict-like access,
+``Axes.spines`` now also supports some ``pandas.Series``-like features.
+
+Accessing single elements by item or by attribute
+
+    ax.spines['top'].set_visible(False)
+    ax.spines.top.set_visible(False)
+
+Accessing a subset of items::
+
+    ax.spines[['top', 'right']].set_visible(False)
+
+Accessing all items simultaneously::
+
+    ax.spines[:].set_visible(False)
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -535,7 +535,7 @@ def __init__(self, fig, rect,
         # placeholder for any colorbars added that use this axes.
         # (see colorbar.py):
         self._colorbars = []
-        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,114 @@ 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 proxy cannot be used for any other operations on its members.
+
+    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 AttributeError(
+                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.
+
+    The interface is dict-like mapping names (e.g. 'left') to `.Spine` objects.
+    Additionally it implements some pandas.Series-like features like accessing
+    elements by attribute::
+
+        spines['top'].set_visible(False)
+        spines.top.set_visible(False)
+
+    Multiple spines can be addressed simultaneously by passing a list. This
+    will return a `SpinesProxy` that broadcasts all ``set_*`` calls to it's
+    members::
+
+        spines[['top', 'right']].set_visible(False)
+
+    Use an open slice to address all spines::
+
+        spines[:].set_visible(False)
+
+    """
+    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 __getattr__(self, name):
+        try:
+            return self._dict[name]
+        except KeyError:
+            raise ValueError(
+                f"'Spines' object does not contain a '{name}' spine")
+
+    def __getitem__(self, key):
+        if isinstance(key, list):
+            unknown_keys = [k for k in key if k not in self._dict]
+            if unknown_keys:
+                raise KeyError(', '.join(unknown_keys))
+            return SpinesProxy({k: v for k, v in self._dict.items()
+                                if k in key})
+        if isinstance(key, tuple):
+            raise ValueError('Multiple spines must be passed as a single list')
+        if isinstance(key, slice):
+            if key.start is None and key.stop is None and key.step is None:
+                return SpinesProxy(self._dict)
+            else:
+                raise ValueError(
+                    'Spines does not support slicing except for the fully '
+                    'open slice [:] to access all spines.')
+        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)
diff --git a/lib/matplotlib/tests/test_spines.py b/lib/matplotlib/tests/test_spines.py
@@ -1,9 +1,52 @@
 import numpy as np
+import pytest
 
 import matplotlib.pyplot as plt
+from matplotlib.spines import Spines
 from matplotlib.testing.decorators import check_figures_equal, image_comparison
 
 
+def test_spine_class():
+    """Test Spines and SpinesProxy in isolation."""
+    class SpineMock:
+        def __init__(self):
+            self.val = None
+
+        def set_val(self, val):
+            self.val = val
+
+    spines_dict = {
+        'left': SpineMock(),
+        'right': SpineMock(),
+        'top': SpineMock(),
+        'bottom': SpineMock(),
+    }
+    spines = Spines(**spines_dict)
+
+    assert spines['left'] is spines_dict['left']
+    assert spines.left is spines_dict['left']
+
+    spines[['left', 'right']].set_val('x')
+    assert spines.left.val == 'x'
+    assert spines.right.val == 'x'
+    assert spines.top.val is None
+    assert spines.bottom.val is None
+
+    spines[:].set_val('y')
+    assert all(spine.val == 'y' for spine in spines.values())
+
+    with pytest.raises(KeyError, match='foo'):
+        spines['foo']
+    with pytest.raises(KeyError, match='foo, bar'):
+        spines[['left', 'foo', 'right', 'bar']]
+    with pytest.raises(ValueError, match='single list'):
+        spines['left', 'right']
+    with pytest.raises(ValueError, match='Spines does not support slicing'):
+        spines['left':'right']
+    with pytest.raises(ValueError, match='Spines does not support slicing'):
+        spines['top':]
+
+
 @image_comparison(['spines_axes_positions'])
 def test_spines_axes_positions():
     # SF bug 2852168
