Add test case for loading figure from a subprocess, and refactor the code.

wsykala · wsykala · commit 2a3de64cac64 · 2022-07-22T14:55:48.000+02:00
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -1,11 +1,13 @@
 from io import BytesIO
+import ast
 import pickle
 
 import numpy as np
 import pytest
 
 import matplotlib as mpl
 from matplotlib import cm
+from matplotlib.testing import subprocess_run_helper
 from matplotlib.testing.decorators import check_figures_equal
 from matplotlib.dates import rrulewrapper
 from matplotlib.lines import VertexSelector
@@ -42,9 +44,7 @@ def test_simple():
     pickle.dump(fig, BytesIO(), pickle.HIGHEST_PROTOCOL)
 
 
-@mpl.style.context("default")
-@check_figures_equal(extensions=["png"])
-def test_complete(fig_test, fig_ref):
+def _generate_complete_test_figure(fig_ref):
     fig_ref.set_size_inches((10, 6))
     plt.figure(fig_ref)
 
@@ -83,12 +83,17 @@ def test_complete(fig_test, fig_ref):
     plt.quiver(x, y, u, v)
 
     plt.subplot(3, 3, 8)
-    plt.scatter(x, x**2, label='$x^2$')
+    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_complete(fig_test, fig_ref):
+    _generate_complete_test_figure(fig_ref)
     # plotting is done, now test its pickle-ability
     pkl = BytesIO()
     pickle.dump(fig_ref, pkl, pickle.HIGHEST_PROTOCOL)
@@ -101,6 +106,46 @@ def test_complete(fig_test, fig_ref):
     plt.close(loaded)
 
 
+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)))
+
+
+@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,
+        extra_env={'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()
