Merge pull request matplotlib#22484 from anntzer/cp

Simplify the internal API to connect picklable callbacks.

Author: QuLogic

lib/matplotlib/axes/_base.py

@@ -658,9 +658,8 @@ def __init__(self, fig, rect,
         self._internal_update(kwargs)
 
         for name, axis in self._axis_map.items():
-            axis.callbacks._pickled_cids.add(
-                axis.callbacks.connect(
-                    'units', self._unit_change_handler(name)))
+            axis.callbacks._connect_picklable(
+                'units', self._unit_change_handler(name))
 
         rcParams = mpl.rcParams
         self.tick_params(

lib/matplotlib/cbook/__init__.py

@@ -236,6 +236,16 @@ def connect(self, signal, func):
         self.callbacks[signal][cid] = proxy
         return cid
 
+    def _connect_picklable(self, signal, func):
+        """
+        Like `.connect`, but the callback is kept when pickling/unpickling.
+
+        Currently internal-use only.
+        """
+        cid = self.connect(signal, func)
+        self._pickled_cids.add(cid)
+        return cid
+
     # Keep a reference to sys.is_finalizing, as sys may have been cleared out
     # at that point.
     def _remove_proxy(self, proxy, *, _is_finalizing=sys.is_finalizing):

lib/matplotlib/tests/test_cbook.py

@@ -183,10 +183,10 @@ def setup(self):
         self.callbacks = cbook.CallbackRegistry()
 
     def connect(self, s, func, pickle):
-        cid = self.callbacks.connect(s, func)
         if pickle:
-            self.callbacks._pickled_cids.add(cid)
-        return cid
+            return self.callbacks.connect(s, func)
+        else:
+            return self.callbacks._connect_picklable(s, func)
 
     def disconnect(self, cid):
         return self.callbacks.disconnect(cid)

lib/mpl_toolkits/mplot3d/axes3d.py

@@ -157,14 +157,12 @@ def __init__(
         self.fmt_zdata = None
 
         self.mouse_init()
-        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.figure.canvas.callbacks._connect_picklable(
+            'motion_notify_event', self._on_move)
+        self.figure.canvas.callbacks._connect_picklable(
+            'button_press_event', self._button_press)
+        self.figure.canvas.callbacks._connect_picklable(
+            'button_release_event', self._button_release)
         self.set_top_view()
 
         self.patch.set_linewidth(0)