Merge pull request #16220 from anntzer/3dp

QuLogic · web-flow · commit 772a315a24aa · 2020-12-15T01:58:32.000-05:00
Fix interaction with unpickled 3d plots.
diff --git a/doc/api/next_api_changes/behavior/16220-AL.rst b/doc/api/next_api_changes/behavior/16220-AL.rst
@@ -0,0 +1,8 @@
+Canvas's callback registry now stored on Figure
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The canonical location of the `~.cbook.CallbackRegistry` used to
+handle Figure/Canvas events has been moved from the Canvas to the
+Figure.  This change should be transparent to almost all users,
+however if you are swapping switching the Figure out from on top of a
+Canvas or visa versa you may see a change in behavior.
diff --git a/lib/matplotlib/backend_bases.py b/lib/matplotlib/backend_bases.py
@@ -1713,8 +1713,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
@@ -1725,6 +1723,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:
@@ -157,21 +167,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
@@ -2322,6 +2322,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)
