Support pickling of figures with aligned x/y labels.

anntzer · anntzer · commit 79381e5dabdf · 2023-02-26T12:17:28.000+01:00
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -788,6 +788,19 @@ class Grouper:
     def __init__(self, init=()):
         self._mapping = {weakref.ref(x): [weakref.ref(x)] for x in init}
 
+    def _as_strong_dict(self):
+        """Export the grouper's weak refs to strong refs; useful for pickling."""
+        self.clean()
+        return {k(): [v() for v in vs] for k, vs in self._mapping.items()}
+
+    @classmethod
+    def _from_strong_dict(cls, mapping):
+        """Create a grouper from a dict of strong refs; useful for unpickling."""
+        self = cls()
+        self._mapping = {weakref.ref(k): [*map(weakref.ref, vs)]
+                         for k, vs in mapping.items()}
+        return self
+
     def __contains__(self, item):
         return weakref.ref(item) in self._mapping
 
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -3174,6 +3174,10 @@ def __getstate__(self):
         from matplotlib import _pylab_helpers
         if self.canvas.manager in _pylab_helpers.Gcf.figs.values():
             state['_restore_to_pylab'] = True
+
+        state["_align_label_groups"] = {
+            k: v._as_strong_dict() for k, v in self._align_label_groups.items()}
+
         return state
 
     def __setstate__(self, state):
@@ -3201,6 +3205,10 @@ def __setstate__(self, state):
             pylab_helpers.Gcf._set_new_active_manager(mgr)
             plt.draw_if_interactive()
 
+        self._align_label_groups = {
+            k: cbook.Grouper._from_strong_dict(v)
+            for k, v in self._align_label_groups.items()}
+
         self.stale = True
 
     def add_axobserver(self, func):
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -58,6 +58,7 @@ def _generate_complete_test_figure(fig_ref):
     # Ensure lists also pickle correctly.
     plt.subplot(3, 3, 1)
     plt.plot(list(range(10)))
+    plt.ylabel("hello")
 
     plt.subplot(3, 3, 2)
     plt.contourf(data, hatches=['//', 'ooo'])
@@ -68,6 +69,7 @@ def _generate_complete_test_figure(fig_ref):
 
     plt.subplot(3, 3, 4)
     plt.imshow(data)
+    plt.ylabel("hello\nworld!")
 
     plt.subplot(3, 3, 5)
     plt.pcolor(data)
@@ -89,6 +91,8 @@ def _generate_complete_test_figure(fig_ref):
     plt.subplot(3, 3, 9)
     plt.errorbar(x, x * -0.5, xerr=0.2, yerr=0.4)
 
+    fig_ref.align_ylabels()
+
 
 @mpl.style.context("default")
 @check_figures_equal(extensions=["png"])
