Fix picklability of make_norm_from_scale norms.

anntzer · anntzer · commit d3ed88a58a6d · 2021-12-17T03:12:53.000+01:00
And also fix their qualname, which was missed in a00a909.
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -1507,9 +1507,26 @@ class norm_cls(Normalize):
 
     if init is None:
         def init(vmin=None, vmax=None, clip=False): pass
-    bound_init_signature = inspect.signature(init)
+
+    return _make_norm_from_scale(
+        scale_cls, base_norm_cls, inspect.signature(init))
+
+
+@functools.lru_cache(None)
+def _make_norm_from_scale(scale_cls, base_norm_cls, bound_init_signature):
+    """
+    Helper for `make_norm_from_scale`.
+
+    This function is split out so that it takes a signature object as third
+    argument (as signatures are picklable, contrary to arbitrary lambdas);
+    caching is also used so that different unpickles reuse the same class.
+    """
 
     class Norm(base_norm_cls):
+        def __reduce__(self):
+            return (_picklable_norm_constructor,
+                    (scale_cls, base_norm_cls, bound_init_signature),
+                    self.__dict__)
 
         def __init__(self, *args, **kwargs):
             ba = bound_init_signature.bind(*args, **kwargs)
@@ -1519,6 +1536,10 @@ def __init__(self, *args, **kwargs):
             self._scale = scale_cls(axis=None, **ba.arguments)
             self._trf = self._scale.get_transform()
 
+        __init__.__signature__ = bound_init_signature.replace(parameters=[
+            inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
+            *bound_init_signature.parameters.values()])
+
         def __call__(self, value, clip=None):
             value, is_scalar = self.process_value(value)
             if self.vmin is None or self.vmax is None:
@@ -1566,17 +1587,23 @@ def autoscale_None(self, A):
             in_trf_domain = np.extract(np.isfinite(self._trf.transform(A)), A)
             return super().autoscale_None(in_trf_domain)
 
-    Norm.__name__ = (f"{scale_cls.__name__}Norm" if base_norm_cls is Normalize
-                     else base_norm_cls.__name__)
-    Norm.__qualname__ = base_norm_cls.__qualname__
+    Norm.__name__ = (
+            f"{scale_cls.__name__}Norm" if base_norm_cls is Normalize
+            else base_norm_cls.__name__)
+    Norm.__qualname__ = (
+            f"{scale_cls.__qualname__}Norm" if base_norm_cls is Normalize
+            else base_norm_cls.__qualname__)
     Norm.__module__ = base_norm_cls.__module__
     Norm.__doc__ = base_norm_cls.__doc__
-    Norm.__init__.__signature__ = bound_init_signature.replace(parameters=[
-        inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
-        *bound_init_signature.parameters.values()])
+
     return Norm
 
 
+def _picklable_norm_constructor(*args, **kwargs):
+    cls = _make_norm_from_scale(*args, **kwargs)
+    return cls.__new__(cls)
+
+
 @make_norm_from_scale(
     scale.FuncScale,
     init=lambda functions, vmin=None, vmax=None, clip=False: None)
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -1480,4 +1480,4 @@ def test_norm_update_figs(fig_test, fig_ref):
 def test_make_norm_from_scale_name():
     logitnorm = mcolors.make_norm_from_scale(
         mscale.LogitScale, mcolors.Normalize)
-    assert logitnorm.__name__ == "LogitScaleNorm"
+    assert logitnorm.__name__ == logitnorm.__qualname__ == "LogitScaleNorm"
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -219,3 +219,10 @@ def test_unpickle_canvas():
 def test_mpl_toolkits():
     ax = parasite_axes.host_axes([0, 0, 1, 1])
     assert type(pickle.loads(pickle.dumps(ax))) == parasite_axes.HostAxes
+
+
+def test_dynamic_norm():
+    logit_norm_instance = mpl.colors.make_norm_from_scale(
+        mpl.scale.LogitScale, mpl.colors.Normalize)()
+    assert type(pickle.loads(pickle.dumps(logit_norm_instance))) \
+        == type(logit_norm_instance)
