From bd0df1f17e5065604a4b9c7da9de3e4c36027675 Mon Sep 17 00:00:00 2001 From: Thomas A Caswell Date: Fri, 29 Jul 2022 10:56:12 -0400 Subject: [PATCH] Backport PR #23462: Fix AttributeError for pickle load of Figure class There were a number of complications with the backport: - the new test generalized an existing test that had picked up additional changes on the main branch - the signature of the subproccess helper has changed on the main branch --- lib/matplotlib/figure.py | 3 +- lib/matplotlib/tests/test_pickle.py | 101 ++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 1 deletion(-) diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py index d114801ea12f..58d071f58e64 100644 --- a/lib/matplotlib/figure.py +++ b/lib/matplotlib/figure.py @@ -2919,7 +2919,8 @@ def __setstate__(self, state): import matplotlib._pylab_helpers as pylab_helpers allnums = plt.get_fignums() num = max(allnums) + 1 if allnums else 1 - mgr = plt._backend_mod.new_figure_manager_given_figure(num, self) + backend = plt._get_backend_mod() + mgr = backend.new_figure_manager_given_figure(num, self) pylab_helpers.Gcf._set_new_active_manager(mgr) plt.draw_if_interactive() diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py index c8b9260bc291..0c499108f8f3 100644 --- a/lib/matplotlib/tests/test_pickle.py +++ b/lib/matplotlib/tests/test_pickle.py @@ -1,4 +1,5 @@ from io import BytesIO +import ast import pickle import numpy as np @@ -6,6 +7,9 @@ from matplotlib import cm from matplotlib.testing.decorators import image_comparison +import matplotlib as mpl +from matplotlib.testing import subprocess_run_helper +from matplotlib.testing.decorators import check_figures_equal from matplotlib.dates import rrulewrapper import matplotlib.pyplot as plt import matplotlib.transforms as mtransforms @@ -110,6 +114,103 @@ def test_complete(): assert fig.get_label() == 'Figure with a label?' +def _pickle_load_subprocess(): + import os + import pickle + + path = os.environ['PICKLE_FILE_PATH'] + + with open(path, 'rb') as blob: + fig = pickle.load(blob) + + print(str(pickle.dumps(fig))) + + +def _generate_complete_test_figure(fig_ref): + fig_ref.set_size_inches((10, 6)) + plt.figure(fig_ref) + + plt.suptitle('Can you fit any more in a figure?') + + # make some arbitrary data + x, y = np.arange(8), np.arange(10) + data = u = v = np.linspace(0, 10, 80).reshape(10, 8) + v = np.sin(v * -0.6) + + # Ensure lists also pickle correctly. + plt.subplot(3, 3, 1) + plt.plot(list(range(10))) + + plt.subplot(3, 3, 2) + plt.contourf(data, hatches=['//', 'ooo']) + plt.colorbar() + + plt.subplot(3, 3, 3) + plt.pcolormesh(data) + + plt.subplot(3, 3, 4) + plt.imshow(data) + + plt.subplot(3, 3, 5) + plt.pcolor(data) + + ax = plt.subplot(3, 3, 6) + ax.set_xlim(0, 7) + ax.set_ylim(0, 9) + plt.streamplot(x, y, u, v) + + ax = plt.subplot(3, 3, 7) + ax.set_xlim(0, 7) + ax.set_ylim(0, 9) + plt.quiver(x, y, u, v) + + plt.subplot(3, 3, 8) + plt.scatter(x, x ** 2, label='$x^2$') + plt.legend(loc='upper left') + + plt.subplot(3, 3, 9) + plt.errorbar(x, x * -0.5, xerr=0.2, yerr=0.4) + + +@mpl.style.context("default") +@check_figures_equal(extensions=['png']) +def test_pickle_load_from_subprocess(fig_test, fig_ref, tmp_path): + _generate_complete_test_figure(fig_ref) + + fp = tmp_path / 'sinus.pickle' + assert not fp.exists() + + with fp.open('wb') as file: + pickle.dump(fig_ref, file, pickle.HIGHEST_PROTOCOL) + assert fp.exists() + + proc = subprocess_run_helper( + _pickle_load_subprocess, + timeout=60, + PICKLE_FILE_PATH=str(fp), + ) + + loaded_fig = pickle.loads(ast.literal_eval(proc.stdout)) + + loaded_fig.canvas.draw() + + fig_test.set_size_inches(loaded_fig.get_size_inches()) + fig_test.figimage(loaded_fig.canvas.renderer.buffer_rgba()) + + plt.close(loaded_fig) + + +def test_gcf(): + fig = plt.figure("a label") + buf = BytesIO() + pickle.dump(fig, buf, pickle.HIGHEST_PROTOCOL) + plt.close("all") + assert plt._pylab_helpers.Gcf.figs == {} # No figures must be left. + fig = pickle.loads(buf.getbuffer()) + assert plt._pylab_helpers.Gcf.figs != {} # A manager is there again. + assert fig.get_label() == "a label" + + def test_no_pyplot(): # tests pickle-ability of a figure not created with pyplot from matplotlib.backends.backend_pdf import FigureCanvasPdf