Merge pull request matplotlib#1175 from pelson/pickle2

Pickling support added. Various whitespace fixes as a result of reading *lots* of code.

Author: efiring

doc/users/whats_new.rst

@@ -100,6 +100,17 @@ minimum and maximum colorbar extensions.
     plt.show()
 
 
+Figures are picklable
+---------------------
+
+Philip Elson added an experimental feature to make figures picklable
+for quick and easy short-term storage of plots. Pickle files
+are not designed for long term storage, are unsupported when restoring a pickle
+saved in another matplotlib version and are insecure when restoring a pickle
+from an untrusted source. Having said this, they are useful for short term
+storage for later modification inside matplotlib.
+
+
 Set default bounding box in matplotlibrc
 ------------------------------------------
 

lib/matplotlib/__init__.py

@@ -1085,6 +1085,7 @@ def tk_window_focus():
     'matplotlib.tests.test_mathtext',
     'matplotlib.tests.test_mlab',
     'matplotlib.tests.test_patches',
+    'matplotlib.tests.test_pickle',
     'matplotlib.tests.test_rcparams',
     'matplotlib.tests.test_simplification',
     'matplotlib.tests.test_spines',

lib/matplotlib/_pylab_helpers.py

@@ -14,7 +14,7 @@ def error_msg(msg):
 
 class Gcf(object):
     """
-    Manage a set of integer-numbered figures.
+    Singleton to manage a set of integer-numbered figures.
 
     This class is never instantiated; it consists of two class
     attributes (a list and a dictionary), and a set of static
@@ -132,6 +132,7 @@ def set_active(manager):
         Gcf._activeQue.append(manager)
         Gcf.figs[manager.num] = manager
 
+
 atexit.register(Gcf.destroy_all)
 
 

lib/matplotlib/artist.py

@@ -103,6 +103,13 @@ def __init__(self):
         self._gid = None
         self._snap = None
 
+    def __getstate__(self):
+        d = self.__dict__.copy()
+        # remove the unpicklable remove method, this will get re-added on load
+        # (by the axes) if the artist lives on an axes.
+        d['_remove_method'] = None
+        return d
+
     def remove(self):
         """
         Remove the artist from the figure if possible.  The effect
@@ -122,7 +129,7 @@ def remove(self):
         # the _remove_method attribute directly.  This would be a protected
         # attribute if Python supported that sort of thing.  The callback
         # has one parameter, which is the child to be removed.
-        if self._remove_method != None:
+        if self._remove_method is not None:
             self._remove_method(self)
         else:
             raise NotImplementedError('cannot remove artist')

lib/matplotlib/axes.py

@@ -153,9 +153,8 @@ def set_default_color_cycle(clist):
                                                     DeprecationWarning)
 
 
-class _process_plot_var_args:
+class _process_plot_var_args(object):
     """
-
     Process variable length arguments to the plot command, so that
     plot commands like the following are supported::
 
@@ -171,6 +170,14 @@ def __init__(self, axes, command='plot'):
         self.command = command
         self.set_color_cycle()
 
+    def __getstate__(self):
+        # note: it is not possible to pickle a itertools.cycle instance
+        return {'axes': self.axes, 'command': self.command}
+
+    def __setstate__(self, state):
+        self.__dict__ = state.copy()
+        self.set_color_cycle()
+
     def set_color_cycle(self, clist=None):
         if clist is None:
             clist = rcParams['axes.color_cycle']
@@ -281,7 +288,7 @@ def _plot_args(self, tup, kwargs):
             linestyle, marker, color = _process_plot_format(tup[-1])
             tup = tup[:-1]
         elif len(tup) == 3:
-            raise ValueError, 'third arg must be a format string'
+            raise ValueError('third arg must be a format string')
         else:
             linestyle, marker, color = None, None, None
         kw = {}
@@ -354,6 +361,7 @@ class Axes(martist.Artist):
 
     def __str__(self):
         return "Axes(%g,%g;%gx%g)" % tuple(self._position.bounds)
+
     def __init__(self, fig, rect,
                  axisbg = None, # defaults to rc axes.facecolor
                  frameon = True,
@@ -488,6 +496,15 @@ def __init__(self, fig, rect,
             self._ycid = self.yaxis.callbacks.connect('units finalize',
                                                       self.relim)
 
+    def __setstate__(self, state):
+        self.__dict__ = state
+        # put the _remove_method back on all artists contained within the axes
+        for container_name in ['lines', 'collections', 'tables', 'patches',
+                               'texts', 'images']:
+            container = getattr(self, container_name)
+            for artist in container:
+                artist._remove_method = container.remove
+
     def get_window_extent(self, *args, **kwargs):
         """
         get the axes bounding box in display space; *args* and
@@ -1472,7 +1489,7 @@ def _update_line_limits(self, line):
             return
 
         line_trans = line.get_transform()
-        
+
         if line_trans == self.transData:
             data_path = path
 
@@ -1491,8 +1508,8 @@ def _update_line_limits(self, line):
             else:
                 data_path = trans_to_data.transform_path(path)
         else:
-            # for backwards compatibility we update the dataLim with the 
-            # coordinate range of the given path, even though the coordinate 
+            # for backwards compatibility we update the dataLim with the
+            # coordinate range of the given path, even though the coordinate
             # systems are completely different. This may occur in situations
             # such as when ax.transAxes is passed through for absolute
             # positioning.
@@ -1502,7 +1519,7 @@ def _update_line_limits(self, line):
             updatex, updatey = line_trans.contains_branch_seperately(
                                                                self.transData
                                                                     )
-            self.dataLim.update_from_path(data_path, 
+            self.dataLim.update_from_path(data_path,
                                           self.ignore_existing_data_limits,
                                           updatex=updatex,
                                           updatey=updatey)
@@ -2199,11 +2216,11 @@ def ticklabel_format(self, **kwargs):
                 cb = False
             else:
                 cb = True
-                raise NotImplementedError, "comma style remains to be added"
+                raise NotImplementedError("comma style remains to be added")
         elif style == '':
             sb = None
         else:
-            raise ValueError, "%s is not a valid style value"
+            raise ValueError("%s is not a valid style value")
         try:
             if sb is not None:
                 if axis == 'both' or axis == 'x':
@@ -3706,9 +3723,9 @@ def hlines(self, y, xmin, xmax, colors='k', linestyles='solid',
             xmax = np.resize( xmax, y.shape )
 
         if len(xmin)!=len(y):
-            raise ValueError, 'xmin and y are unequal sized sequences'
+            raise ValueError('xmin and y are unequal sized sequences')
         if len(xmax)!=len(y):
-            raise ValueError, 'xmax and y are unequal sized sequences'
+            raise ValueError('xmax and y are unequal sized sequences')
 
         verts = [ ((thisxmin, thisy), (thisxmax, thisy))
                             for thisxmin, thisxmax, thisy in zip(xmin, xmax, y)]
@@ -3785,9 +3802,9 @@ def vlines(self, x, ymin, ymax, colors='k', linestyles='solid',
             ymax = np.resize( ymax, x.shape )
 
         if len(ymin)!=len(x):
-            raise ValueError, 'ymin and x are unequal sized sequences'
+            raise ValueError('ymin and x are unequal sized sequences')
         if len(ymax)!=len(x):
-            raise ValueError, 'ymax and x are unequal sized sequences'
+            raise ValueError('ymax and x are unequal sized sequences')
 
         Y = np.array([ymin, ymax]).T
 
@@ -4768,7 +4785,7 @@ def make_iterable(x):
             if len(height) == 1:
                 height *= nbars
         else:
-            raise ValueError, 'invalid orientation: %s' % orientation
+            raise ValueError('invalid orientation: %s' % orientation)
 
         if len(linewidth) < nbars:
             linewidth *= nbars
@@ -4826,7 +4843,7 @@ def make_iterable(x):
                 bottom = [bottom[i] - height[i]/2. for i in xrange(len(bottom))]
 
         else:
-            raise ValueError, 'invalid alignment: %s' % align
+            raise ValueError('invalid alignment: %s' % align)
 
         args = zip(left, bottom, width, height, color, edgecolor, linewidth)
         for l, b, w, h, c, e, lw in args:
@@ -5701,7 +5718,7 @@ def computeConfInterval(data, med, iq, bootstrap):
                 else:
                     x = [x[:,i] for i in xrange(nc)]
             else:
-                raise ValueError, "input x can have no more than 2 dimensions"
+                raise ValueError("input x can have no more than 2 dimensions")
         if not hasattr(x[0], '__len__'):
             x = [x]
         col = len(x)
@@ -7069,10 +7086,10 @@ def _pcolorargs(self, funcname, *args):
 
         Nx = X.shape[-1]
         Ny = Y.shape[0]
-        if len(X.shape) <> 2 or X.shape[0] == 1:
+        if len(X.shape) != 2 or X.shape[0] == 1:
             x = X.reshape(1,Nx)
             X = x.repeat(Ny, axis=0)
-        if len(Y.shape) <> 2 or Y.shape[1] == 1:
+        if len(Y.shape) != 2 or Y.shape[1] == 1:
             y = Y.reshape(Ny, 1)
             Y = y.repeat(Nx, axis=1)
         if X.shape != Y.shape:
@@ -8815,7 +8832,15 @@ def __init__(self, fig, *args, **kwargs):
         # _axes_class is set in the subplot_class_factory
         self._axes_class.__init__(self, fig, self.figbox, **kwargs)
 
-
+    def __reduce__(self):
+        # get the first axes class which does not inherit from a subplotbase
+        not_subplotbase = lambda c: issubclass(c, Axes) and \
+                                    not issubclass(c, SubplotBase)
+        axes_class = [c for c in self.__class__.mro() if not_subplotbase(c)][0]
+        r = [_PicklableSubplotClassConstructor(),
+             (axes_class,),
+             self.__getstate__()]
+        return tuple(r)
 
     def get_geometry(self):
         """get the subplot geometry, eg 2,2,3"""
@@ -8897,6 +8922,22 @@ def subplot_class_factory(axes_class=None):
 # This is provided for backward compatibility
 Subplot = subplot_class_factory()
 
+
+class _PicklableSubplotClassConstructor(object):
+    """
+    This stub class exists to return the appropriate subplot
+    class when __call__-ed with an axes class. This is purely to
+    allow Pickling of Axes and Subplots.
+    """
+    def __call__(self, axes_class):
+        # create a dummy object instance
+        subplot_instance = _PicklableSubplotClassConstructor()
+        subplot_class = subplot_class_factory(axes_class)
+        # update the class to the desired subplot class
+        subplot_instance.__class__ = subplot_class
+        return subplot_instance
+
+
 docstring.interpd.update(Axes=martist.kwdoc(Axes))
 docstring.interpd.update(Subplot=martist.kwdoc(Axes))
 

lib/matplotlib/axis.py

@@ -597,7 +597,6 @@ class Ticker:
     formatter = None
 
 
-
 class Axis(artist.Artist):
 
     """

lib/matplotlib/backends/__init__.py

@@ -52,6 +52,6 @@ def do_nothing(*args, **kwargs): pass
 
     matplotlib.verbose.report('backend %s version %s' % (backend,backend_version))
 
-    return new_figure_manager, draw_if_interactive, show
+    return backend_mod, new_figure_manager, draw_if_interactive, show
 
 

lib/matplotlib/backends/backend_agg.py

@@ -385,7 +385,6 @@ def post_processing(image, dpi):
                                       image)
 
 
-
 def new_figure_manager(num, *args, **kwargs):
     """
     Create a new figure manager instance
@@ -396,7 +395,14 @@ def new_figure_manager(num, *args, **kwargs):
 
     FigureClass = kwargs.pop('FigureClass', Figure)
     thisFig = FigureClass(*args, **kwargs)
-    canvas = FigureCanvasAgg(thisFig)
+    return new_figure_manager_given_figure(num, thisFig)
+
+
+def new_figure_manager_given_figure(num, figure):
+    """
+    Create a new figure manager instance for the given figure.
+    """
+    canvas = FigureCanvasAgg(figure)
     manager = FigureManagerBase(canvas, num)
     return manager