Merge pull request #4915 from QuLogic/transformwrapper-pickles

mdboom · tacaswell · commit cba1283ea14e · 2015-10-08T12:47:55.000-04:00
TransformWrapper pickling fixes
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -12,6 +12,7 @@
 
 from matplotlib.testing.decorators import cleanup, image_comparison
 import matplotlib.pyplot as plt
+import matplotlib.transforms as mtransforms
 
 
 def depth_getter(obj,
@@ -252,6 +253,34 @@ def test_polar():
     plt.draw()
 
 
+class TransformBlob(object):
+    def __init__(self):
+        self.identity = mtransforms.IdentityTransform()
+        self.identity2 = mtransforms.IdentityTransform()
+        # Force use of the more complex composition.
+        self.composite = mtransforms.CompositeGenericTransform(
+            self.identity,
+            self.identity2)
+        # Check parent -> child links of TransformWrapper.
+        self.wrapper = mtransforms.TransformWrapper(self.composite)
+        # Check child -> parent links of TransformWrapper.
+        self.composite2 = mtransforms.CompositeGenericTransform(
+            self.wrapper,
+            self.identity)
+
+
+def test_transform():
+    obj = TransformBlob()
+    pf = pickle.dumps(obj)
+    del obj
+
+    obj = pickle.loads(pf)
+    # Check parent -> child links of TransformWrapper.
+    assert_equal(obj.wrapper._child, obj.composite)
+    # Check child -> parent links of TransformWrapper.
+    assert_equal(list(obj.wrapper._parents.values()), [obj.composite2])
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=['-s'])
diff --git a/lib/matplotlib/transforms.py b/lib/matplotlib/transforms.py
@@ -1534,6 +1534,10 @@ def __init__(self, child):
             msg = ("'child' must be an instance of"
                    " 'matplotlib.transform.Transform'")
             raise ValueError(msg)
+        self._init(child)
+        self.set_children(child)
+
+    def _init(self, child):
         Transform.__init__(self)
         self.input_dims = child.input_dims
         self.output_dims = child.output_dims
@@ -1549,12 +1553,18 @@ def __str__(self):
             return str(self._child)
 
     def __getstate__(self):
-        # only store the child
-        return {'child': self._child}
+        # only store the child and parents
+        return {
+            'child': self._child,
+            # turn the weakkey dictionary into a normal dictionary
+            'parents': dict(six.iteritems(self._parents))
+        }
 
     def __setstate__(self, state):
         # re-initialise the TransformWrapper with the state's child
-        self.__init__(state['child'])
+        self._init(state['child'])
+        # turn the normal dictionary back into a WeakValueDictionary
+        self._parents = WeakValueDictionary(state['parents'])
 
     def __repr__(self):
         return "TransformWrapper(%r)" % self._child
@@ -1565,7 +1575,6 @@ def frozen(self):
 
     def _set(self, child):
         self._child = child
-        self.set_children(child)
 
         self.transform = child.transform
         self.transform_affine = child.transform_affine
@@ -1594,6 +1603,7 @@ def set(self, child):
                    " output dimensions as the current child.")
             raise ValueError(msg)
 
+        self.set_children(child)
         self._set(child)
 
         self._invalid = 0
