FIX: Allow deepcopy on norms and scales

jklymak · jklymak · commit e152f0204f0c · 2021-01-09T17:39:06.000-08:00
diff --git a/doc/api/scale_api.rst b/doc/api/scale_api.rst
@@ -3,6 +3,7 @@
 ********************
 
 .. automodule:: matplotlib.scale
+   :private-members: _CopyableTransformMixin
    :members:
    :undoc-members:
    :show-inheritance:
diff --git a/lib/matplotlib/scale.py b/lib/matplotlib/scale.py
@@ -79,6 +79,17 @@ def limit_range_for_scale(self, vmin, vmax, minpos):
         return vmin, vmax
 
 
+class _CopyableTransformMixin():
+    """
+    Mixin to support copy and deep copy on transforms.  This alows scales,
+    and hence norms, to be copyable.
+    """
+    def __deepcopy__(self, memo):
+        return self.frozen()
+
+    __copy__ = __deepcopy__
+
+
 class LinearScale(ScaleBase):
     """
     The default linear scale.
@@ -113,9 +124,9 @@ def get_transform(self):
         return IdentityTransform()
 
 
-class FuncTransform(Transform):
+class FuncTransform(_CopyableTransformMixin, Transform):
     """
-    A simple transform that takes and arbitrary function for the
+    A transform that takes and arbitrary function for the
     forward and inverse transform.
     """
 
@@ -191,7 +202,7 @@ def set_default_locators_and_formatters(self, axis):
             axis.set_minor_locator(NullLocator())
 
 
-class LogTransform(Transform):
+class LogTransform(_CopyableTransformMixin, Transform):
     input_dims = output_dims = 1
 
     @_api.rename_parameter("3.3", "nonpos", "nonpositive")
@@ -233,7 +244,7 @@ def inverted(self):
         return InvertedLogTransform(self.base)
 
 
-class InvertedLogTransform(Transform):
+class InvertedLogTransform(_CopyableTransformMixin, Transform):
     input_dims = output_dims = 1
 
     def __init__(self, base):
@@ -359,7 +370,7 @@ def get_transform(self):
         return self._transform
 
 
-class SymmetricalLogTransform(Transform):
+class SymmetricalLogTransform(_CopyableTransformMixin, Transform):
     input_dims = output_dims = 1
 
     def __init__(self, base, linthresh, linscale):
@@ -391,7 +402,7 @@ def inverted(self):
                                                self.linscale)
 
 
-class InvertedSymmetricalLogTransform(Transform):
+class InvertedSymmetricalLogTransform(_CopyableTransformMixin, Transform):
     input_dims = output_dims = 1
 
     def __init__(self, base, linthresh, linscale):
@@ -494,7 +505,7 @@ def get_transform(self):
         return self._transform
 
 
-class LogitTransform(Transform):
+class LogitTransform(_CopyableTransformMixin, Transform):
     input_dims = output_dims = 1
 
     @_api.rename_parameter("3.3", "nonpos", "nonpositive")
@@ -520,7 +531,7 @@ def __str__(self):
         return "{}({!r})".format(type(self).__name__, self._nonpositive)
 
 
-class LogisticTransform(Transform):
+class LogisticTransform(_CopyableTransformMixin, Transform):
     input_dims = output_dims = 1
 
     @_api.rename_parameter("3.3", "nonpos", "nonpositive")
diff --git a/lib/matplotlib/tests/test_colors.py b/lib/matplotlib/tests/test_colors.py
@@ -16,6 +16,7 @@
 import matplotlib.colorbar as mcolorbar
 import matplotlib.cbook as cbook
 import matplotlib.pyplot as plt
+import matplotlib.scale as mscale
 from matplotlib.testing.decorators import image_comparison
 
 
@@ -1320,3 +1321,16 @@ def test_2d_to_rgba():
     rgba_1d = mcolors.to_rgba(color.reshape(-1))
     rgba_2d = mcolors.to_rgba(color.reshape((1, -1)))
     assert rgba_1d == rgba_2d
+
+
+def test_norm_deepcopy():
+    norm = mcolors.LogNorm()
+    norm.vmin = 0.0002
+    norm2 = copy.deepcopy(norm)
+    assert norm2.vmin == norm.vmin
+    assert isinstance(norm2._scale, mscale.LogScale)
+    norm = mcolors.Normalize()
+    norm.vmin = 0.0002
+    norm2 = copy.deepcopy(norm)
+    assert isinstance(norm2._scale, mscale.LinearScale)
+    assert norm2.vmin == norm.vmin
diff --git a/lib/matplotlib/tests/test_scale.py b/lib/matplotlib/tests/test_scale.py
@@ -7,6 +7,7 @@
 
 import numpy as np
 from numpy.testing import assert_allclose
+import copy
 import io
 import pytest
 
@@ -210,3 +211,9 @@ def test_pass_scale():
     ax.set_yscale(scale)
     assert ax.xaxis.get_scale() == 'log'
     assert ax.yaxis.get_scale() == 'log'
+
+
+def test_scale_deepcopy():
+    sc = mscale.LogScale(axis='x', base=10)
+    sc2 = copy.deepcopy(sc)
+    assert str(sc.get_transform()) == str(sc2.get_transform())
diff --git a/lib/matplotlib/transforms.py b/lib/matplotlib/transforms.py
@@ -2062,6 +2062,12 @@ def frozen(self):
         # docstring inherited
         return self
 
+    def __deepcopy__(self, memo):
+        # The identity transform does not need to lock out deepcopy
+        return self
+
+    __copy__ = __deepcopy__
+
     __str__ = _make_str_method()
 
     def get_matrix(self):
