Add ability to record pickled figure output

Dumps the current figure into a pickle file when ``saveas`` ends
with ".pickle".

CHANGES.rst

@@ -22,7 +22,12 @@ Release History
 1.0.2 (unreleased)
 ==================
 
+**Added**
 
+- Added the ability to save plots as pickle files using the
+  ``.pickle`` extension with ``plt.saveas``. (`#23`_)
+
+.. _#23: https://github.com/nengo/pytest-plt/pull/23
 
 1.0.1 (October 28, 2019)
 ========================

README.rst

@@ -99,6 +99,18 @@ if the PDF format is unsuitable.
 
    plt.saveas = "%s.png" % (plt.saveas[:-4],)
 
+Moreover, using the extension ``.pickle`` will tell pytest-plt to pickle the
+current figure object. The figure can then be inspected using pyplot's
+interactive GUI after unpickling the file. You can achieve this with the
+following code snippet.
+
+.. code-block:: python
+
+   import pickle
+   import matplotlib.pyplot as plt
+   pickle.load(open('path/to/my/plot/figure.pickle', 'rb'))
+   plt.show()
+
 Configuration
 =============
 

pytest_plt/plugin.py

@@ -2,6 +2,7 @@
 
 import errno
 import os
+import pickle
 import re
 
 from matplotlib import use as mpl_use
@@ -147,10 +148,16 @@ def __exit__(self, type, value, traceback):
 
     def save(self, path):
         mkdir_p(os.path.dirname(path))
-        savefig_kw = {"bbox_inches": "tight"}
-        if hasattr(self.plt, "bbox_extra_artists"):
-            savefig_kw["bbox_extra_artists"] = self.plt.bbox_extra_artists
-        self.plt.savefig(path, **savefig_kw)
+
+        if path.endswith(".pickle"):
+            with open(path, "wb") as fh:
+                pickle.dump(self.plt.gcf(), fh)
+        else:
+            savefig_kw = {"bbox_inches": "tight"}
+            if hasattr(self.plt, "bbox_extra_artists"):
+                savefig_kw["bbox_extra_artists"] = self.plt.bbox_extra_artists
+            self.plt.savefig(path, **savefig_kw)
+
         super(Plotter, self).save(path)
 
 

pytest_plt/tests/test_plt.py

@@ -42,3 +42,11 @@ def test_bbox_extra_artists(plt):
 def test_saveas(plt):
     assert plt.saveas.endswith("saveas.pdf")
     plt.saveas = None
+
+
+def test_saveas_pickle(plt):
+    axes = plt.subplots(2, 3)[1]  # The pickled figure will contain six axes.
+    x = np.linspace(-1, 1, 21)
+    for k, ax in enumerate(axes.ravel()):
+        ax.plot(x, x ** k)
+    plt.saveas = "%s.pickle" % (plt.saveas[:-4],)

pytest_plt/tests/test_pytest.py

@@ -12,7 +12,9 @@
 test files can be run manually by passing them to ``pytest``.
 """
 
+import os
 from pathlib import Path
+import pickle
 
 import pytest
 from pytest_plt.plugin import Mock
@@ -145,7 +147,7 @@ def test_filename_drop_prefix(testdir, prefix):
         path = Path(plot)
         assert path.parts[0] == "plots"
         assert path.stem == plot_name
-        assert path.suffix in [".pdf", ".png"]
+        assert path.suffix in [".pdf", ".png", ".pickle"]
         assert path.exists()
 
 
@@ -221,3 +223,27 @@ def test_default_dir(testdir):
         assert path.parts[0] == "myoverridedir"
         assert path.name.startswith("package.tests.")
         assert path.exists()
+
+
+def test_pickle_files(testdir):
+    """Verify that pickle files can be loaded and contain the correct figure.
+
+    The figure is checked by the number of axes (see test_plt.py::test_saveas_pickle).
+    It is also written to the "plots" folder in `testdir` as a ".png".
+    """
+    copy_all_tests(testdir, "package/tests")
+
+    result = testdir.runpytest("-v", "--plots")
+    saved_files = [Path(plot) for _, plot in saved_plots(result)]
+
+    saved_pickle_files = [path for path in saved_files if path.suffix == ".pickle"]
+    for pickle_file in saved_pickle_files:
+        fig = pickle.load(open(pickle_file, "rb"))
+        fig.savefig(os.path.splitext(pickle_file)[0] + ".png")
+        assert len(fig.axes) == 6
+
+    # verify that other output file formats are not mistakenly being pickled
+    saved_img_files = [path for path in saved_files if path.suffix != ".pickle"]
+    for img_file in saved_img_files:
+        with pytest.raises(pickle.UnpicklingError):
+            pickle.load(open(img_file, "rb"))