Restore axes sharedness when unpickling.

anntzer · anntzer · commit 859a412be003 · 2018-07-01T09:31:56.000+02:00
Previously, pickling and unpickling shared axes would result in
axes sharing a ticker instance (because that's how shared axes
are set up), but without changes of one's xlims propagated to the
other.  The reason is that that sharedness information is stored in
AxesBase._shared_x_axes, which does *not* get pickled together with the
Axes instance: the latter only has a textual reference "I am an instance
of AxesBase", so the Grouper information is lost.

To keep the Grouper information valid, additionally pickle the current
group of shared axes together with the Axes, and restore that
information upon unpickling.
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -542,17 +542,25 @@ def __getstate__(self):
         # The renderer should be re-created by the figure, and then cached at
         # that point.
         state = super().__getstate__()
-        state['_cachedRenderer'] = None
-        state.pop('_layoutbox')
-        state.pop('_poslayoutbox')
-
+        for key in ['_cachedRenderer', '_layoutbox', '_poslayoutbox']:
+            state[key] = None
+        # Prune the sharing & twinning info to only contain the current group.
+        for grouper_name in [
+                '_shared_x_axes', '_shared_y_axes', '_twinned_axes']:
+            grouper = getattr(self, grouper_name)
+            state[grouper_name] = (grouper.get_siblings(self)
+                                   if self in grouper else None)
         return state
 
     def __setstate__(self, state):
+        # Merge the grouping info back into the global groupers.
+        for grouper_name in [
+                '_shared_x_axes', '_shared_y_axes', '_twinned_axes']:
+            siblings = state.pop(grouper_name)
+            if siblings:
+                getattr(self, grouper_name).join(*siblings)
         self.__dict__ = state
         self._stale = True
-        self._layoutbox = None
-        self._poslayoutbox = None
 
     def get_window_extent(self, *args, **kwargs):
         """
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -185,3 +185,10 @@ def test_rrulewrapper():
     except RecursionError:
         print('rrulewrapper pickling test failed')
         raise
+
+
+def test_shared():
+    fig, axs = plt.subplots(2, sharex=True)
+    fig = pickle.loads(pickle.dumps(fig))
+    fig.axes[0].set_xlim(10, 20)
+    assert fig.axes[1].get_xlim() == (10, 20)
