Make inset_axes and secondary_axis picklable.

anntzer · anntzer · commit 623ed984138e · 2020-07-21T12:35:49.000+02:00
... by making the locator a class, instead of a nested function.
Also, reuse the inset_locator machinery in secondary_axis, instead of
duplicating it there.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -37,31 +37,29 @@
 _log = logging.getLogger(__name__)
 
 
-def _make_inset_locator(bounds, trans, parent):
+class _InsetLocator:
     """
-    Helper function to locate inset axes, used in
-    `.Axes.inset_axes`.
+    Axes locator for `.Axes.inset_axes`.
 
-    A locator gets used in `Axes.set_aspect` to override the default
-    locations...  It is a function that takes an axes object and
-    a renderer and tells `set_aspect` where it is to be placed.
-
-    Here *rect* is a rectangle [l, b, w, h] that specifies the
-    location for the axes in the transform given by *trans* on the
-    *parent*.
+    The locater is a callable object used in `.Axes.set_aspect` to compute the
+    axes location depending on the renderer.
     """
-    _bounds = mtransforms.Bbox.from_bounds(*bounds)
-    _trans = trans
-    _parent = parent
 
-    def inset_locator(ax, renderer):
-        bbox = _bounds
-        bb = mtransforms.TransformedBbox(bbox, _trans)
-        tr = _parent.figure.transFigure.inverted()
-        bb = mtransforms.TransformedBbox(bb, tr)
-        return bb
+    def __init__(self, bounds, transform):
+        """
+        *bounds* (a ``[l, b, w, h]`` rectangle) and *transform* together
+        specify the position of the inset axes.
+        """
+        self._bounds = bounds
+        self._transform = transform
 
-    return inset_locator
+    def __call__(self, ax, renderer):
+        # Subtracting transFigure will typically rely on inverted(), freezing
+        # the transform; thus, this needs to be delayed until draw time as
+        # transFigure may otherwise change after this is evaluated.
+        return mtransforms.TransformedBbox(
+            mtransforms.Bbox.from_bounds(*self._bounds),
+            self._transform - ax.figure.transFigure)
 
 
 # The axes module contains all the wrappers to plotting functions.
@@ -468,10 +466,9 @@ def inset_axes(self, bounds, *, transform=None, zorder=5, **kwargs):
         kwargs.setdefault('label', 'inset_axes')
 
         # This puts the rectangle into figure-relative coordinates.
-        inset_locator = _make_inset_locator(bounds, transform, self)
-        bb = inset_locator(None, None)
-
-        inset_ax = Axes(self.figure, bb.bounds, zorder=zorder, **kwargs)
+        inset_locator = _InsetLocator(bounds, transform)
+        bounds = inset_locator(self, None).bounds
+        inset_ax = Axes(self.figure, bounds, zorder=zorder, **kwargs)
         # this locator lets the axes move if in data coordinates.
         # it gets called in `ax.apply_aspect() (of all places)
         inset_ax.set_axes_locator(inset_locator)
diff --git a/lib/matplotlib/axes/_secondary_axes.py b/lib/matplotlib/axes/_secondary_axes.py
@@ -4,36 +4,10 @@
 import matplotlib.docstring as docstring
 import matplotlib.ticker as mticker
 import matplotlib.transforms as mtransforms
+from matplotlib.axes import _axes
 from matplotlib.axes._base import _AxesBase
 
 
-def _make_secondary_locator(rect, parent):
-    """
-    Helper function to locate the secondary axes.
-
-    A locator gets used in `Axes.set_aspect` to override the default
-    locations...  It is a function that takes an axes object and
-    a renderer and tells `set_aspect` where it is to be placed.
-
-    This locator make the transform be in axes-relative co-coordinates
-    because that is how we specify the "location" of the secondary axes.
-
-    Here *rect* is a rectangle [l, b, w, h] that specifies the
-    location for the axes in the transform given by *trans* on the
-    *parent*.
-    """
-    _rect = mtransforms.Bbox.from_bounds(*rect)
-    def secondary_locator(ax, renderer):
-        # delay evaluating transform until draw time because the
-        # parent transform may have changed (i.e. if window reesized)
-        bb = mtransforms.TransformedBbox(_rect, parent.transAxes)
-        tr = parent.figure.transFigure.inverted()
-        bb = mtransforms.TransformedBbox(bb, tr)
-        return bb
-
-    return secondary_locator
-
-
 class SecondaryAxis(_AxesBase):
     """
     General class to hold a Secondary_X/Yaxis.
@@ -137,11 +111,13 @@ def set_location(self, location):
         self._loc = location
 
         if self._orientation == 'x':
+            # An x-secondary axes is like an inset axes from x = 0 to x = 1 and
+            # from y = pos to y = pos + eps, in the parent's transAxes coords.
             bounds = [0, self._pos, 1., 1e-10]
         else:
             bounds = [self._pos, 0, 1e-10, 1]
 
-        secondary_locator = _make_secondary_locator(bounds, self._parent)
+        secondary_locator = _axes._InsetLocator(bounds, self._parent.transAxes)
 
         # this locator lets the axes move in the parent axes coordinates.
         # so it never needs to know where the parent is explicitly in
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -193,6 +193,13 @@ def test_shared():
     assert fig.axes[1].get_xlim() == (10, 20)
 
 
+def test_inset_and_secondary():
+    fig, ax = plt.subplots()
+    ax.inset_axes([.1, .1, .3, .3])
+    ax.secondary_xaxis("top", functions=(np.square, np.sqrt))
+    pickle.loads(pickle.dumps(fig))
+
+
 @pytest.mark.parametrize("cmap", cm._cmap_registry.values())
 def test_cmap(cmap):
     pickle.dumps(cmap)
