Templatize class factories.

anntzer · anntzer · commit 8d1a0315ff7a · 2021-05-09T11:40:59.000+02:00
This makes mpl_toolkits axes classes picklable (see test_pickle) by
generalizing the machinery of _picklable_subplot_class_constructor,
which would otherwise have had to be reimplemented for each class
factory.
diff --git a/doc/api/next_api_changes/removals/19033-AL.rst b/doc/api/next_api_changes/removals/19033-AL.rst
@@ -0,0 +1,2 @@
+The private ``matplotlib.axes._subplots._subplot_classes`` dict has been removed
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/lib/matplotlib/axes/_subplots.py b/lib/matplotlib/axes/_subplots.py
@@ -1,5 +1,3 @@
-import functools
-
 from matplotlib import _api, docstring
 import matplotlib.artist as martist
 from matplotlib.axes._axes import Axes
@@ -37,15 +35,6 @@ def __init__(self, fig, *args, **kwargs):
         # This will also update the axes position.
         self.set_subplotspec(SubplotSpec._from_subplot_args(fig, args))
 
-    def __reduce__(self):
-        # get the first axes class which does not inherit from a subplotbase
-        axes_class = next(
-            c for c in type(self).__mro__
-            if issubclass(c, Axes) and not issubclass(c, SubplotBase))
-        return (_picklable_subplot_class_constructor,
-                (axes_class,),
-                self.__getstate__())
-
     @_api.deprecated(
         "3.4", alternative="get_subplotspec",
         addendum="(get_subplotspec returns a SubplotSpec instance.)")
@@ -170,59 +159,10 @@ def _make_twin_axes(self, *args, **kwargs):
         return twin
 
 
-# this here to support cartopy which was using a private part of the
-# API to register their Axes subclasses.
-
-# In 3.1 this should be changed to a dict subclass that warns on use
-# In 3.3 to a dict subclass that raises a useful exception on use
-# In 3.4 should be removed
-
-# The slow timeline is to give cartopy enough time to get several
-# release out before we break them.
-_subplot_classes = {}
-
-
-@functools.lru_cache(None)
-def subplot_class_factory(axes_class=None):
-    """
-    Make a new class that inherits from `.SubplotBase` and the
-    given axes_class (which is assumed to be a subclass of `.axes.Axes`).
-    This is perhaps a little bit roundabout to make a new class on
-    the fly like this, but it means that a new Subplot class does
-    not have to be created for every type of Axes.
-    """
-    if axes_class is None:
-        _api.warn_deprecated(
-            "3.3", message="Support for passing None to subplot_class_factory "
-            "is deprecated since %(since)s; explicitly pass the default Axes "
-            "class instead. This will become an error %(removal)s.")
-        axes_class = Axes
-    try:
-        # Avoid creating two different instances of GeoAxesSubplot...
-        # Only a temporary backcompat fix.  This should be removed in
-        # 3.4
-        return next(cls for cls in SubplotBase.__subclasses__()
-                    if cls.__bases__ == (SubplotBase, axes_class))
-    except StopIteration:
-        return type("%sSubplot" % axes_class.__name__,
-                    (SubplotBase, axes_class),
-                    {'_axes_class': axes_class})
-
-
+subplot_class_factory = cbook._make_class_factory(
+    SubplotBase, "{}Subplot", "_axes_class", default_axes_class=Axes)
 Subplot = subplot_class_factory(Axes)  # Provided for backward compatibility.
 
-
-def _picklable_subplot_class_constructor(axes_class):
-    """
-    Stub factory that returns an empty instance of the appropriate subplot
-    class when called with an axes class. This is purely to allow pickling of
-    Axes and Subplots.
-    """
-    subplot_class = subplot_class_factory(axes_class)
-    return subplot_class.__new__(subplot_class)
-
-
 docstring.interpd.update(Axes_kwdoc=martist.kwdoc(Axes))
-docstring.dedent_interpd(Axes.__init__)
-
 docstring.interpd.update(Subplot_kwdoc=martist.kwdoc(Axes))
+docstring.dedent_interpd(Axes.__init__)
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -2298,3 +2298,54 @@ def _unikey_or_keysym_to_mplkey(unikey, keysym):
         "next": "pagedown",  # Used by tk.
     }.get(key, key)
     return key
+
+
+@functools.lru_cache(None)
+def _make_class_factory(
+        mixin_class, fmt, axes_attr=None, *, default_axes_class=None):
+    """
+    Return a function that creates picklable classes inheriting from a mixin.
+
+    After ::
+
+        factory = _make_class_factory(FooMixin, fmt, axes_attr)
+        FooAxes = factory(Axes)
+
+    ``Foo`` is a class that inherits from ``FooMixin`` and ``Axes`` and **is
+    picklable** (picklability is what differentiates this from a plain call to
+    `type`).  Its ``__name__`` is set to ``fmt.format(Axes.__name__)`` and the
+    base class is stored in the ``axes_attr`` attribute, if not None.
+    """
+
+    @functools.lru_cache(None)
+    def class_factory(axes_class):
+        # default_axes_class should go away once the deprecation elapses.
+        if axes_class is None and default_axes_class is not None:
+            warn_deprecated(
+                "3.3", message="Support for passing None to class factories "
+                "is deprecated since %(since)s and will be removed "
+                "%(removal)s; explicitly pass the default Axes class instead.")
+            return class_factory(default_axes_class)
+        d = {"__reduce__":
+             lambda self: (_picklable_class_constructor,
+                           (mixin_class, fmt, axes_attr,
+                            default_axes_class, axes_class,),
+                           self.__getstate__())}
+        if axes_attr is not None:
+            d[axes_attr] = axes_class
+        cls = type(
+            fmt.format(axes_class.__name__), (mixin_class, axes_class), d)
+        # Better in first approximation than __module__ = "matplotlib.cbook"...
+        cls.__module__ = mixin_class.__module__
+        return cls
+
+    return class_factory
+
+
+def _picklable_class_constructor(
+        base_cls, fmt, axes_attr, default_axes_class, axes_class):
+    """Internal helper for _make_class_factory."""
+    cls = _make_class_factory(
+        base_cls, fmt, axes_attr,
+        default_axes_class=default_axes_class)(axes_class)
+    return cls.__new__(cls)
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -6325,21 +6325,6 @@ def test_spines_properbbox_after_zoom():
     np.testing.assert_allclose(bb.get_points(), bb2.get_points(), rtol=1e-6)
 
 
-def test_cartopy_backcompat():
-
-    class Dummy(matplotlib.axes.Axes):
-        ...
-
-    class DummySubplot(matplotlib.axes.SubplotBase, Dummy):
-        _axes_class = Dummy
-
-    matplotlib.axes._subplots._subplot_classes[Dummy] = DummySubplot
-
-    FactoryDummySubplot = matplotlib.axes.subplot_class_factory(Dummy)
-
-    assert DummySubplot is FactoryDummySubplot
-
-
 def test_gettightbbox_ignore_nan():
     fig, ax = plt.subplots()
     remove_ticks_and_titles(fig)
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -10,6 +10,7 @@
 import matplotlib.pyplot as plt
 import matplotlib.transforms as mtransforms
 import matplotlib.figure as mfigure
+from mpl_toolkits.axes_grid1 import parasite_axes
 
 
 def test_simple():
@@ -212,3 +213,8 @@ def test_unpickle_canvas():
     out.seek(0)
     fig2 = pickle.load(out)
     assert fig2.canvas is not None
+
+
+def test_mpl_toolkits():
+    ax = parasite_axes.host_axes([0, 0, 1, 1])
+    assert type(pickle.loads(pickle.dumps(ax))) == parasite_axes.HostAxes
diff --git a/lib/mpl_toolkits/axes_grid1/parasite_axes.py b/lib/mpl_toolkits/axes_grid1/parasite_axes.py
@@ -94,12 +94,8 @@ def apply_aspect(self, position=None):
     # end of aux_transform support
 
 
-@functools.lru_cache(None)
-def parasite_axes_class_factory(axes_class):
-    return type("%sParasite" % axes_class.__name__,
-                (ParasiteAxesBase, axes_class), {})
-
-
+parasite_axes_class_factory = cbook._make_class_factory(
+    ParasiteAxesBase, "{}Parasite")
 ParasiteAxes = parasite_axes_class_factory(Axes)
 
 
@@ -283,7 +279,7 @@ def _add_twin_axes(self, axes_class, **kwargs):
         *kwargs* are forwarded to the parasite axes constructor.
         """
         if axes_class is None:
-            axes_class = self._get_base_axes()
+            axes_class = self._base_axes_class
         ax = parasite_axes_class_factory(axes_class)(self, **kwargs)
         self.parasites.append(ax)
         ax._remove_method = self._remove_any_twin
@@ -310,11 +306,10 @@ def get_tightbbox(self, renderer, call_axes_locator=True,
         return Bbox.union([b for b in bbs if b.width != 0 or b.height != 0])
 
 
-@functools.lru_cache(None)
-def host_axes_class_factory(axes_class):
-    return type("%sHostAxes" % axes_class.__name__,
-                (HostAxesBase, axes_class),
-                {'_get_base_axes': lambda self: axes_class})
+host_axes_class_factory = cbook._make_class_factory(
+    HostAxesBase, "{}HostAxes", "_base_axes_class")
+HostAxes = host_axes_class_factory(Axes)
+SubplotHost = subplot_class_factory(HostAxes)
 
 
 def host_subplot_class_factory(axes_class):
@@ -323,10 +318,6 @@ def host_subplot_class_factory(axes_class):
     return subplot_host_class
 
 
-HostAxes = host_axes_class_factory(Axes)
-SubplotHost = subplot_class_factory(HostAxes)
-
-
 def host_axes(*args, axes_class=Axes, figure=None, **kwargs):
     """
     Create axes that can act as a hosts to parasitic axes.
diff --git a/lib/mpl_toolkits/axisartist/floating_axes.py b/lib/mpl_toolkits/axisartist/floating_axes.py
@@ -5,10 +5,9 @@
 # TODO :
 # see if tick_iterator method can be simplified by reusing the parent method.
 
-import functools
-
 import numpy as np
 
+from matplotlib import cbook
 import matplotlib.patches as mpatches
 from matplotlib.path import Path
 import matplotlib.axes as maxes
@@ -360,13 +359,8 @@ def adjust_axes_lim(self):
         self.set_ylim(ymin-dy, ymax+dy)
 
 
-@functools.lru_cache(None)
-def floatingaxes_class_factory(axes_class):
-    return type("Floating %s" % axes_class.__name__,
-                (FloatingAxesBase, axes_class),
-                {'_axes_class_floating': axes_class})
-
-
+floatingaxes_class_factory = cbook._make_class_factory(
+    FloatingAxesBase, "Floating {}", "_axes_class_floating")
 FloatingAxes = floatingaxes_class_factory(
     host_axes_class_factory(axislines.Axes))
 FloatingSubplot = maxes.subplot_class_factory(FloatingAxes)

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	+The private ``matplotlib.axes._subplots._subplot_classes`` dict has been removed
	`2`	`+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~`