Accept .pkl extension for pickled files

I've always used .pkl for pickled files in the past.

I also added some more detail to the example in the documentation
and made it into its own section, and fixed two static check errors
resulting from noqa comments being on the wrong lines.

CHANGES.rst

@@ -25,7 +25,7 @@ Release History
 **Added**
 
 - Added the ability to save plots as pickle files using the
-  ``.pickle`` extension with ``plt.saveas``. (`#23`_)
+  ``.pkl`` or ``.pickle`` extensions with ``plt.saveas``. (`#23`_)
 
 .. _#23: https://github.com/nengo/pytest-plt/pull/23
 

README.rst

@@ -99,17 +99,39 @@ 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.
+Using plt.show
+--------------
+
+If you want to inspect a figure using
+the GUI that pops up with ``plt.show()``,
+you can achieve this by saving a plot with the
+``.pkl`` or ``.pickle`` extension.
+pytest-plt will pickle the current figure object
+for later inspection.
+
+Building on the previous example,
+you can change the ``saveas`` attribute like so:
+
+.. code-block:: python
+
+   def test_rectification(plt):
+       values = list(range(-10, 11))
+       rectified = [v if v > 0 else 0 for v in values]
+       assert all(v >= 0 for v in rectified)
+       plt.plot(values, label="Original")
+       plt.plot(rectified, label="Rectified")
+       plt.legend()
+       plt.saveas = "test_rec.pkl"
+
+Then use the following snippet to inspect the figure.
 
 .. code-block:: python
 
-   import pickle
-   import matplotlib.pyplot as plt
-   pickle.load(open('path/to/my/plot/figure.pickle', 'rb'))
-   plt.show()
+   >>> import pickle
+   >>> import matplotlib.pyplot as plt
+   >>> with open("path/to/test_rec.pkl", "rb") as fh:
+   ...     fig = pickle.load(fh)
+   >>> plt.show()
 
 Configuration
 =============

pytest_plt/plugin.py

@@ -7,9 +7,9 @@
 
 from matplotlib import use as mpl_use
 
-mpl_use("Agg")  # noqa: E402
-from matplotlib import pyplot as mpl_plt
-import pytest
+mpl_use("Agg")
+from matplotlib import pyplot as mpl_plt  # noqa: E402
+import pytest  # noqa: E402
 
 
 def mkdir_p(path):
@@ -149,7 +149,7 @@ def __exit__(self, type, value, traceback):
     def save(self, path):
         mkdir_p(os.path.dirname(path))
 
-        if path.endswith(".pickle"):
+        if path.endswith(".pkl") or path.endswith(".pickle"):
             with open(path, "wb") as fh:
                 pickle.dump(self.plt.gcf(), fh)
         else:

pytest_plt/tests/test_plt.py

@@ -3,6 +3,7 @@
 """Test the plt fixture."""
 
 import numpy as np
+import pytest
 
 
 def test_rectification(plt):
@@ -44,9 +45,10 @@ def test_saveas(plt):
     plt.saveas = None
 
 
-def test_saveas_pickle(plt):
+@pytest.mark.parametrize("ext", ["pkl", "pickle"])
+def test_saveas_pickle(plt, ext):
     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],)
+    plt.saveas = "%s.%s" % (plt.saveas[:-4], ext)

pytest_plt/tests/test_pytest.py

@@ -147,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", ".pickle"]
+        assert path.suffix in [".pdf", ".png", ".pkl", ".pickle"]
         assert path.exists()
 
 
@@ -236,14 +236,20 @@ def test_pickle_files(testdir):
     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"]
+    saved_pickle_files = [
+        path for path in saved_files if path.suffix in [".pickle", ".pkl"]
+    ]
     for pickle_file in saved_pickle_files:
-        fig = pickle.load(open(pickle_file, "rb"))
-        fig.savefig(os.path.splitext(pickle_file)[0] + ".png")
+        with open(str(pickle_file), "rb") as fh:
+            fig = pickle.load(fh)
+        fig.savefig("%s.png" % (os.path.splitext(pickle_file)[0],))
         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"]
+    saved_img_files = [
+        path for path in saved_files if path.suffix not in [".pkl", ".pickle"]
+    ]
     for img_file in saved_img_files:
         with pytest.raises(pickle.UnpicklingError):
-            pickle.load(open(img_file, "rb"))
+            with open(str(img_file), "rb") as fh:
+                pickle.load(fh)