Merge pull request #1 from joniemi/save_pickled_figures

Save pickled figures

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,15 @@ 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"):
+            pickle.dump(self.plt.gcf(), open(path, "wb"))
+        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,8 @@ def test_bbox_extra_artists(plt):
 def test_saveas(plt):
     assert plt.saveas.endswith("saveas.pdf")
     plt.saveas = None
+
+
+def test_saveas_pickle(plt):
+    plt.subplots(2, 3)  # The pickled figure will contain six axes.
+    plt.saveas = "%s.pickle" % (plt.saveas[:-4],)

pytest_plt/tests/test_pytest.py

@@ -12,6 +12,8 @@
 test files can be run manually by passing them to ``pytest``.
 """
 
+import pickle
+import matplotlib.pyplot as plt
 from pathlib import Path
 
 import pytest
@@ -49,7 +51,7 @@ def assert_all_passed(result):
     """
     outcomes = result.parseoutcomes()
     for outcome in outcomes:
-        if outcome not in ("passed", "seconds", "warnings"):
+        if outcome not in ("passed", "seconds", "warning"):
             assert outcomes[outcome] == 0
     return outcomes.get("passed", 0)
 
@@ -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,41 @@ def test_default_dir(testdir):
         assert path.parts[0] == "myoverridedir"
         assert path.name.startswith("package.tests.")
         assert path.exists()
+
+
+def test_pickle_files_contain_a_figure(testdir):
+    """ Verify that pickle files can be loaded and contain the correct
+    figure. Th figure is tested by simply checking the number of axes.
+
+    For the expected number of axes, see test_plt.py::test_saveas_pickle
+
+    """
+    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:
+        try:
+            pickle.load(open(pickle_file, "rb"))
+            assert 6 == len(plt.gcf().axes)
+        except pickle.UnpicklingError:
+            assert False, "Could not read a pickled file {}".format(str(pickle_file))
+
+
+def test_image_files_are_not_pickled(testdir):
+    """ Verify that other output file formats are not mistakenly being
+    pickled.
+    """
+    copy_all_tests(testdir, "package/tests")
+
+    result = testdir.runpytest("-v", "--plots")
+
+    saved_files = [Path(plot) for _, plot in saved_plots(result)]
+    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"))