|
1 | 1 | fromioimportBytesIO |
| 2 | +importast |
2 | 3 | importpickle |
3 | 4 |
|
4 | 5 | importnumpyasnp |
5 | 6 | importpytest |
6 | 7 |
|
7 | 8 | frommatplotlibimportcm |
8 | 9 | frommatplotlib.testing.decoratorsimportimage_comparison |
| 10 | +importmatplotlibasmpl |
| 11 | +frommatplotlib.testingimportsubprocess_run_helper |
| 12 | +frommatplotlib.testing.decoratorsimportcheck_figures_equal |
9 | 13 | frommatplotlib.datesimportrrulewrapper |
10 | 14 | importmatplotlib.pyplotasplt |
11 | 15 | importmatplotlib.transformsasmtransforms |
@@ -110,6 +114,103 @@ def test_complete(): |
110 | 114 | assertfig.get_label()=='Figure with a label?' |
111 | 115 |
|
112 | 116 |
|
| 117 | +def_pickle_load_subprocess(): |
| 118 | +importos |
| 119 | +importpickle |
| 120 | + |
| 121 | +path=os.environ['PICKLE_FILE_PATH'] |
| 122 | + |
| 123 | +withopen(path,'rb')asblob: |
| 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 | +deftest_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 | +assertnotfp.exists() |
| 182 | + |
| 183 | +withfp.open('wb')asfile: |
| 184 | +pickle.dump(fig_ref,file,pickle.HIGHEST_PROTOCOL) |
| 185 | +assertfp.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 | +deftest_gcf(): |
| 204 | +fig=plt.figure("a label") |
| 205 | +buf=BytesIO() |
| 206 | +pickle.dump(fig,buf,pickle.HIGHEST_PROTOCOL) |
| 207 | +plt.close("all") |
| 208 | +assertplt._pylab_helpers.Gcf.figs== {}# No figures must be left. |
| 209 | +fig=pickle.loads(buf.getbuffer()) |
| 210 | +assertplt._pylab_helpers.Gcf.figs!= {}# A manager is there again. |
| 211 | +assertfig.get_label()=="a label" |
| 212 | + |
| 213 | + |
113 | 214 | deftest_no_pyplot(): |
114 | 215 | # tests pickle-ability of a figure not created with pyplot |
115 | 216 | frommatplotlib.backends.backend_pdfimportFigureCanvasPdf |
|