Merge pull request #3627 from pelson/pickle_bug

tacaswell · tacaswell · commit 97eb61202fbc · 2014-10-10T12:49:44.000-04:00
BUG : Fixed Image and Renderer pickling
diff --git a/lib/matplotlib/backends/backend_agg.py b/lib/matplotlib/backends/backend_agg.py
@@ -104,6 +104,14 @@ def __init__(self, width, height, dpi):
         if __debug__: verbose.report('RendererAgg.__init__ done',
                                      'debug-annoying')
 
+    def __getstate__(self):
+        # We only want to preserve the init keywords of the Renderer.
+        # Anything else can be re-created.
+        return {'width': self.width, 'height': self.height, 'dpi': self.dpi}
+
+    def __setstate__(self, state):
+        self.__init__(state['width'], state['height'], state['dpi'])
+
     def _get_hinting_flag(self):
         if rcParams['text.hinting']:
             return LOAD_FORCE_AUTOHINT
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -10,10 +10,8 @@
 
 import os
 import warnings
-import math
 
 import numpy as np
-from numpy import ma
 
 from matplotlib import rcParams
 import matplotlib.artist as martist
@@ -113,6 +111,12 @@ def __init__(self, ax,
 
         self.update(kwargs)
 
+    def __getstate__(self):
+        state = super(_AxesImageBase, self).__getstate__()
+        # We can't pickle the C Image cached object.
+        state.pop('_imcache', None)
+        return state
+
     def get_size(self):
         """Get the numrows, numcols of the input image"""
         if self._A is None:
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -352,6 +352,12 @@ def __init__(self, xdata, ydata,
         self._invalidy = True
         self.set_data(xdata, ydata)
 
+    def __getstate__(self):
+        state = super(Line2D, self).__getstate__()
+        # _linefunc will be restored on draw time.
+        state.pop('_lineFunc', None)
+        return state
+
     def contains(self, mouseevent):
         """
         Test whether the mouse event occurred on the line.  The pick
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -107,6 +107,8 @@ def test_simple():
     plt.plot(list(xrange(10)), label='foobar')
     plt.legend()
 
+    # Uncomment to debug any unpicklable objects. This is slow so is not
+    # uncommented by default.
 #    recursive_pickle(fig)
     pickle.dump(ax, BytesIO(), pickle.HIGHEST_PROTOCOL)
 
@@ -195,17 +197,47 @@ def test_complete():
 
 def test_no_pyplot():
     # tests pickle-ability of a figure not created with pyplot
-
-    import pickle as p
     from matplotlib.backends.backend_pdf import FigureCanvasPdf as fc
     from matplotlib.figure import Figure
 
     fig = Figure()
-    can = fc(fig)
+    _ = fc(fig)
     ax = fig.add_subplot(1, 1, 1)
     ax.plot([1, 2, 3], [1, 2, 3])
-
-    # Uncomment to debug any unpicklable objects. This is slow so is not
-    # uncommented by default.
-#    recursive_pickle(fig)
     pickle.dump(fig, BytesIO(), pickle.HIGHEST_PROTOCOL)
+
+
+def test_renderer():
+    from matplotlib.backends.backend_agg import RendererAgg
+    renderer = RendererAgg(10, 20, 30)
+    pickle.dump(renderer, BytesIO())
+
+
+def test_image():
+    # Prior to v1.4.0 the Image would cache data which was not picklable
+    # once it had been drawn.
+    from matplotlib.backends.backend_agg import new_figure_manager
+    manager = new_figure_manager(1000)
+    fig = manager.canvas.figure
+    ax = fig.add_subplot(1, 1, 1)
+    ax.imshow(np.arange(12).reshape(3, 4))
+    manager.canvas.draw()
+    pickle.dump(fig, BytesIO())
+
+
+def test_grid():
+    from matplotlib.backends.backend_agg import new_figure_manager
+    manager = new_figure_manager(1000)
+    fig = manager.canvas.figure
+    ax = fig.add_subplot(1, 1, 1)
+    ax.grid()
+    # Drawing the grid triggers instance methods to be attached
+    # to the Line2D object (_lineFunc).
+    manager.canvas.draw()
+
+    pickle.dump(ax, BytesIO())
+
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule(argv=['-s'])
