From c47a0637a97f17146dae202baf74a0f687757457 Mon Sep 17 00:00:00 2001 From: Trevor Bekolay Date: Thu, 23 Jul 2020 14:00:32 -0500 Subject: [PATCH] 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 | 2 +- README.rst | 38 ++++++++++++++++++++++++++------- pytest_plt/plugin.py | 8 +++---- pytest_plt/tests/test_plt.py | 6 ++++-- pytest_plt/tests/test_pytest.py | 18 ++++++++++------ 5 files changed, 51 insertions(+), 21 deletions(-) diff --git a/CHANGES.rst b/CHANGES.rst index 15b19d1..b13eaf5 100644 --- a/CHANGES.rst +++ b/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 diff --git a/README.rst b/README.rst index b039553..2e706a7 100644 --- a/README.rst +++ b/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 ============= diff --git a/pytest_plt/plugin.py b/pytest_plt/plugin.py index 3810d5a..d36ddce 100644 --- a/pytest_plt/plugin.py +++ b/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: diff --git a/pytest_plt/tests/test_plt.py b/pytest_plt/tests/test_plt.py index 08d0978..2e7d87b 100644 --- a/pytest_plt/tests/test_plt.py +++ b/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) diff --git a/pytest_plt/tests/test_pytest.py b/pytest_plt/tests/test_pytest.py index f0d95c6..ddbaab6 100644 --- a/pytest_plt/tests/test_pytest.py +++ b/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)