Make all transforms copiable (and thus scales, too).

anntzer · anntzer · commit e62d61ecd88f · 2021-01-12T10:10:08.000+01:00
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
@@ -1,3 +1,5 @@
+import copy
+
 import matplotlib.pyplot as plt
 from matplotlib.scale import (
     LogTransform, InvertedLogTransform,
@@ -210,3 +212,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
@@ -33,6 +33,7 @@
 # `np.minimum` instead of the builtin `min`, and likewise for `max`.  This is
 # done so that `nan`s are propagated, instead of being silently dropped.
 
+import copy
 import functools
 import textwrap
 import weakref
@@ -139,11 +140,26 @@ def __setstate__(self, data_dict):
             k: weakref.ref(v, lambda _, pop=self._parents.pop, k=k: pop(k))
             for k, v in self._parents.items() if v is not None}
 
-    def __copy__(self, *args):
-        raise NotImplementedError(
-            "TransformNode instances can not be copied. "
-            "Consider using frozen() instead.")
-    __deepcopy__ = __copy__
+    def __copy__(self):
+        other = copy.copy(super())
+        # If `c = a + b; a1 = copy(a)`, then modifications to `a1` do not
+        # propagate back to `c`, i.e. we need to clear the parents of `a1`.
+        other._parents = {}
+        # If `c = a + b; c1 = copy(c)`, then modifications to `a` also need to
+        # be propagated to `c1`.
+        for key, val in vars(self).items():
+            if isinstance(val, TransformNode) and id(self) in val._parents:
+                other.set_children(val)  # val == getattr(other, key)
+        return other
+
+    def __deepcopy__(self, memo):
+        other = copy.deepcopy(super(), memo)
+        # If `c = a + b; a1 = copy(a)`, then modifications to `a1` do not
+        # propagate back to `c`, i.e. we need to clear the parents of `a1`.
+        other._parents = {}
+        # If `c = a + b; c1 = copy(c)`, this creates a separate tree
+        # (`c1 = a1 + b1`) so nothing needs to be done.
+        return other
 
     def invalidate(self):
         """
