Fix interaction with unpickled 3d plots.

anntzer · anntzer · commit 6d79e6eddfe1 · 2020-10-28T12:03:17.000+01:00
In order to reset the mouse interaction callbacks on unpickled 3d plots,
the callback registry really needs to be on the Figure object rather
than the Canvas, because the canvas doesn't exist yet when the 3d axes
is being unpickled (it is only set on the figure at the very end of
unpickling).

So move the callback registry to the figure (with a proxy property on
the canvas).

Then, add a private mechanism to pickle select callbacks, and combine
everything together.

Test with e.g.
```
import matplotlib.pyplot as plt
import pickle

fig = plt.figure()
fig.add_subplot(111, projection='3d')

p = pickle.dumps(fig)
plt.close("all")

pickle.loads(p)

plt.show()
```
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -1718,8 +1718,6 @@ def __init__(self, figure):
         figure.set_canvas(self)
         self.figure = figure
         self.manager = None
-        # a dictionary from event name to a dictionary that maps cid->func
-        self.callbacks = cbook.CallbackRegistry()
         self.widgetlock = widgets.LockDraw()
         self._button = None  # the button pressed
         self._key = None  # the key pressed
@@ -1730,6 +1728,10 @@ def __init__(self, figure):
         self.toolbar = None  # NavigationToolbar2 will set me
         self._is_idle_drawing = False
 
+    @property
+    def callbacks(self):
+        return self.figure._canvas_callbacks
+
     @classmethod
     @functools.lru_cache()
     def _fix_ipython_backend2gui(cls):
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -104,6 +104,16 @@ def __hash__(self):
         return hash(self._obj)
 
 
+def _weak_or_strong_ref(func, callback):
+    """
+    Return a `WeakMethod` wrapping *func* if possible, else a `_StrongRef`.
+    """
+    try:
+        return weakref.WeakMethod(func, callback)
+    except TypeError:
+        return _StrongRef(func)
+
+
 class CallbackRegistry:
     """
     Handle registering and disconnecting for a set of signals and callbacks:
@@ -163,21 +173,37 @@ def __init__(self, exception_handler=_exception_printer):
         self.callbacks = {}
         self._cid_gen = itertools.count()
         self._func_cid_map = {}
+        # A hidden variable that marks cids that need to be pickled.
+        self._pickled_cids = set()
 
     def __getstate__(self):
-        # In general, callbacks may not be pickled, so we just drop them.
-        return {**vars(self), "callbacks": {}, "_func_cid_map": {}}
+        return {
+            **vars(self),
+            # In general, callbacks may not be pickled, so we just drop them,
+            # unless directed otherwise by self._pickled_cids.
+            "callbacks": {s: {cid: proxy() for cid, proxy in d.items()
+                              if cid in self._pickled_cids}
+                          for s, d in self.callbacks.items()},
+            # It is simpler to reconstruct this from callbacks in __setstate__.
+            "_func_cid_map": None,
+        }
+
+    def __setstate__(self, state):
+        vars(self).update(state)
+        self.callbacks = {
+            s: {cid: _weak_or_strong_ref(func, self._remove_proxy)
+                for cid, func in d.items()}
+            for s, d in self.callbacks.items()}
+        self._func_cid_map = {
+            s: {proxy: cid for cid, proxy in d.items()}
+            for s, d in self.callbacks.items()}
 
     def connect(self, s, func):
         """Register *func* to be called when signal *s* is generated."""
         self._func_cid_map.setdefault(s, {})
-        try:
-            proxy = weakref.WeakMethod(func, self._remove_proxy)
-        except TypeError:
-            proxy = _StrongRef(func)
+        proxy = _weak_or_strong_ref(func, self._remove_proxy)
         if proxy in self._func_cid_map[s]:
             return self._func_cid_map[s][proxy]
-
         cid = next(self._cid_gen)
         self._func_cid_map[s][proxy] = cid
         self.callbacks.setdefault(s, {})
diff --git a/lib/matplotlib/figure.py b/lib/matplotlib/figure.py
@@ -2280,6 +2280,10 @@ def __init__(self,
         super().__init__()
 
         self.callbacks = cbook.CallbackRegistry()
+        # Callbacks traditionally associated with the canvas (and exposed with
+        # a proxy property), but that actually need to be on the figure for
+        # pickling.
+        self._canvas_callbacks = cbook.CallbackRegistry()
 
         if figsize is None:
             figsize = mpl.rcParams['figure.figsize']
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -118,12 +118,14 @@ def __init__(
             self._zcid = None
 
         self.mouse_init()
-        self.figure.canvas.mpl_connect(
-            'motion_notify_event', self._on_move),
-        self.figure.canvas.mpl_connect(
-            'button_press_event', self._button_press),
-        self.figure.canvas.mpl_connect(
-            'button_release_event', self._button_release),
+        self.figure.canvas.callbacks._pickled_cids.update({
+            self.figure.canvas.mpl_connect(
+                'motion_notify_event', self._on_move),
+            self.figure.canvas.mpl_connect(
+                'button_press_event', self._button_press),
+            self.figure.canvas.mpl_connect(
+                'button_release_event', self._button_release),
+        })
         self.set_top_view()
 
         self.patch.set_linewidth(0)
