Templatize class factories.

anntzer · anntzer · commit cdecca39aaae · 2021-05-24T22:07:24.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,6 @@
+The private ``matplotlib.axes._subplots._subplot_classes`` dict has been removed
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Support for passing ``None`` to ``subplot_class_factory`` is deprecated
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Explicitly pass in the base `~matplotlib.axes.Axes` class instead.
diff --git a/lib/matplotlib/axes/_subplots.py b/lib/matplotlib/axes/_subplots.py
@@ -1,6 +1,4 @@
-import functools
-
-from matplotlib import _api
+from matplotlib import _api, cbook
 from matplotlib.axes._axes import Axes
 from matplotlib.gridspec import GridSpec, SubplotSpec
 
@@ -36,15 +34,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.)")
@@ -169,53 +158,6 @@ 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")
 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)
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -2188,3 +2188,53 @@ 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, attr_name=None):
+    """
+    Return a function that creates picklable classes inheriting from a mixin.
+
+    After ::
+
+        factory = _make_class_factory(FooMixin, fmt, attr_name)
+        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 ``attr_name`` attribute, if not None.
+
+    Moreover, the return value of ``factory`` is memoized: calls with the same
+    ``Axes`` class always return the same subclass.
+    """
+
+    @functools.lru_cache(None)
+    def class_factory(axes_class):
+        # The parameter is named "axes_class" for backcompat but is really just
+        # a base class; no axes semantics are used.
+        base_class = axes_class
+
+        class subcls(mixin_class, base_class):
+            # Better approximation than __module__ = "matplotlib.cbook".
+            __module__ = mixin_class.__module__
+
+            def __reduce__(self):
+                return (_picklable_class_constructor,
+                        (mixin_class, fmt, attr_name, base_class),
+                        self.__getstate__())
+
+        subcls.__name__ = subcls.__qualname__ = fmt.format(base_class.__name__)
+        if attr_name is not None:
+            setattr(subcls, attr_name, base_class)
+        return subcls
+
+    class_factory.__module__ = mixin_class.__module__
+    return class_factory
+
+
+def _picklable_class_constructor(mixin_class, fmt, attr_name, base_class):
+    """Internal helper for _make_class_factory."""
+    factory = _make_class_factory(mixin_class, fmt, attr_name)
+    cls = factory(base_class)
+    return cls.__new__(cls)
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -6340,21 +6340,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
@@ -1,6 +1,6 @@
 import functools
 
-from matplotlib import _api
+from matplotlib import _api, cbook
 import matplotlib.artist as martist
 import matplotlib.image as mimage
 import matplotlib.transforms as mtransforms
@@ -95,12 +95,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)
 
 
@@ -277,7 +273,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
@@ -304,11 +300,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):
@@ -317,10 +312,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
@@ -351,12 +350,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), {})
-
-
+floatingaxes_class_factory = cbook._make_class_factory(
+    FloatingAxesBase, "Floating {}")
 FloatingAxes = floatingaxes_class_factory(
     host_axes_class_factory(axislines.Axes))
 FloatingSubplot = maxes.subplot_class_factory(FloatingAxes)
