|
1 | 1 | from io import BytesIO
|
| 2 | +import ast |
2 | 3 | import pickle
|
3 | 4 |
|
4 | 5 | import numpy as np
|
5 | 6 | import pytest
|
6 | 7 |
|
7 | 8 | from matplotlib import cm
|
8 | 9 | from matplotlib.testing.decorators import image_comparison
|
| 10 | +import matplotlib as mpl |
| 11 | +from matplotlib.testing import subprocess_run_helper |
| 12 | +from matplotlib.testing.decorators import check_figures_equal |
9 | 13 | from matplotlib.dates import rrulewrapper
|
10 | 14 | import matplotlib.pyplot as plt
|
11 | 15 | import matplotlib.transforms as mtransforms
|
@@ -110,6 +114,103 @@ def test_complete():
|
110 | 114 | assert fig.get_label() == 'Figure with a label?'
|
111 | 115 |
|
112 | 116 |
|
| 117 | +def _pickle_load_subprocess(): |
| 118 | + import os |
| 119 | + import pickle |
| 120 | + |
| 121 | + path = os.environ['PICKLE_FILE_PATH'] |
| 122 | + |
| 123 | + with open(path, 'rb') as blob: |
| 124 | + fig = pickle.load(blob) |
| 125 | + |
| 126 | + print(str(pickle.dumps(fig))) |
| 127 | + |
| 128 | + |
| 129 | +def _generate_complete_test_figure(fig_ref): |
| 130 | + fig_ref.set_size_inches((10, 6)) |
| 131 | + plt.figure(fig_ref) |
| 132 | + |
| 133 | + plt.suptitle('Can you fit any more in a figure?') |
| 134 | + |
| 135 | + # make some arbitrary data |
| 136 | + x, y = np.arange(8), np.arange(10) |
| 137 | + data = u = v = np.linspace(0, 10, 80).reshape(10, 8) |
| 138 | + v = np.sin(v * -0.6) |
| 139 | + |
| 140 | + # Ensure lists also pickle correctly. |
| 141 | + plt.subplot(3, 3, 1) |
| 142 | + plt.plot(list(range(10))) |
| 143 | + |
| 144 | + plt.subplot(3, 3, 2) |
| 145 | + plt.contourf(data, hatches=['//', 'ooo']) |
| 146 | + plt.colorbar() |
| 147 | + |
| 148 | + plt.subplot(3, 3, 3) |
| 149 | + plt.pcolormesh(data) |
| 150 | + |
| 151 | + plt.subplot(3, 3, 4) |
| 152 | + plt.imshow(data) |
| 153 | + |
| 154 | + plt.subplot(3, 3, 5) |
| 155 | + plt.pcolor(data) |
| 156 | + |
| 157 | + ax = plt.subplot(3, 3, 6) |
| 158 | + ax.set_xlim(0, 7) |
| 159 | + ax.set_ylim(0, 9) |
| 160 | + plt.streamplot(x, y, u, v) |
| 161 | + |
| 162 | + ax = plt.subplot(3, 3, 7) |
| 163 | + ax.set_xlim(0, 7) |
| 164 | + ax.set_ylim(0, 9) |
| 165 | + plt.quiver(x, y, u, v) |
| 166 | + |
| 167 | + plt.subplot(3, 3, 8) |
| 168 | + plt.scatter(x, x ** 2, label='$x^2$') |
| 169 | + plt.legend(loc='upper left') |
| 170 | + |
| 171 | + plt.subplot(3, 3, 9) |
| 172 | + plt.errorbar(x, x * -0.5, xerr=0.2, yerr=0.4) |
| 173 | + |
| 174 | + |
| 175 | +@mpl.style.context("default") |
| 176 | +@check_figures_equal(extensions=['png']) |
| 177 | +def test_pickle_load_from_subprocess(fig_test, fig_ref, tmp_path): |
| 178 | + _generate_complete_test_figure(fig_ref) |
| 179 | + |
| 180 | + fp = tmp_path / 'sinus.pickle' |
| 181 | + assert not fp.exists() |
| 182 | + |
| 183 | + with fp.open('wb') as file: |
| 184 | + pickle.dump(fig_ref, file, pickle.HIGHEST_PROTOCOL) |
| 185 | + assert fp.exists() |
| 186 | + |
| 187 | + proc = subprocess_run_helper( |
| 188 | + _pickle_load_subprocess, |
| 189 | + timeout=60, |
| 190 | + PICKLE_FILE_PATH=str(fp), |
| 191 | + ) |
| 192 | + |
| 193 | + loaded_fig = pickle.loads(ast.literal_eval(proc.stdout)) |
| 194 | + |
| 195 | + loaded_fig.canvas.draw() |
| 196 | + |
| 197 | + fig_test.set_size_inches(loaded_fig.get_size_inches()) |
| 198 | + fig_test.figimage(loaded_fig.canvas.renderer.buffer_rgba()) |
| 199 | + |
| 200 | + plt.close(loaded_fig) |
| 201 | + |
| 202 | + |
| 203 | +def test_gcf(): |
| 204 | + fig = plt.figure("a label") |
| 205 | + buf = BytesIO() |
| 206 | + pickle.dump(fig, buf, pickle.HIGHEST_PROTOCOL) |
| 207 | + plt.close("all") |
| 208 | + assert plt._pylab_helpers.Gcf.figs == {} # No figures must be left. |
| 209 | + fig = pickle.loads(buf.getbuffer()) |
| 210 | + assert plt._pylab_helpers.Gcf.figs != {} # A manager is there again. |
| 211 | + assert fig.get_label() == "a label" |
| 212 | + |
| 213 | + |
113 | 214 | def test_no_pyplot():
|
114 | 215 | # tests pickle-ability of a figure not created with pyplot
|
115 | 216 | from matplotlib.backends.backend_pdf import FigureCanvasPdf
|
|
0 commit comments