Restore axes sharedness when unpickling.

anntzer · anntzer · commit 970cb575ee32 · 2018-03-03T00:35:59.000-08: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, instead move the Groupers to the
instance dictionaries (as references to global groupers).  Also make
Groupers picklable following a similar strategy as Transforms, i.e. by
transforming weakrefs into real refs when pickling and transforming them
back into weakref when unpickling.
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -406,15 +406,16 @@ def _grab_next_args(self, *args, **kwargs):
             yield from self._plot_args(this, kwargs)
 
 
+_shared_x_axes = cbook.Grouper()
+_shared_y_axes = cbook.Grouper()
+_twinned_axes = cbook.Grouper()
+
+
 class _AxesBase(martist.Artist):
     """
     """
     name = "rectilinear"
 
-    _shared_x_axes = cbook.Grouper()
-    _shared_y_axes = cbook.Grouper()
-    _twinned_axes = cbook.Grouper()
-
     def __str__(self):
         return "{0}({1[0]:g},{1[1]:g};{1[2]:g}x{1[3]:g})".format(
             type(self).__name__, self._position.bounds)
@@ -479,6 +480,13 @@ def __init__(self, fig, rect,
         """ % {'scale': ' | '.join(
             [repr(x) for x in mscale.get_scale_names()])}
         martist.Artist.__init__(self)
+
+        # Reference the global instances in the instance dict to support
+        # pickling.
+        self._shared_x_axes = _shared_x_axes
+        self._shared_y_axes = _shared_y_axes
+        self._twinned_axes = _twinned_axes
+
         if isinstance(rect, mtransforms.Bbox):
             self._position = rect
         else:
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -1479,9 +1479,26 @@ class Grouper(object):
 
     """
     def __init__(self, init=()):
-        mapping = self._mapping = {}
-        for x in init:
-            mapping[ref(x)] = [ref(x)]
+        self._mapping = {ref(x): [ref(x)] for x in init}
+
+    def __getstate__(self):
+        mapping = {}
+        for k, vs in self._mapping.items():
+            k = k()
+            if k is None:
+                continue
+            mapping[k] = l = []
+            for v in vs:
+                v = v()
+                if v is None:
+                    continue
+                l.append(v)
+        return {"_mapping": mapping}
+
+    def __setstate__(self, state):
+        self.__dict__ = state
+        self._mapping = {ref(k): [ref(v) for v in vs]
+                         for k, vs in self._mapping.items()}
 
     def __contains__(self, item):
         return ref(item) in self._mapping
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -189,3 +189,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)
