Simplify the internal API to connect picklable callbacks.

The API may be made public in the future, but right now the point is
only to make its internal use more ergonomic as it can be helpful in
some other internal places too (to be updated in a separate patch).

Author: anntzer

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)