From 4caa5ba85c11cc499eb1a45372eb249173e6b1fd Mon Sep 17 00:00:00 2001
From: Antony Lee <anntzer.lee@gmail.com>
Date: Thu, 17 Feb 2022 22:34:28 +0100
Subject: [PATCH] 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).
---
 lib/matplotlib/axes/_base.py       |  5 ++---
 lib/matplotlib/cbook/__init__.py   | 10 ++++++++++
 lib/matplotlib/tests/test_cbook.py |  6 +++---
 lib/mpl_toolkits/mplot3d/axes3d.py | 14 ++++++--------
 4 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
index f3e5b206a52a..aa53fe1ee03f 100644
--- a/lib/matplotlib/axes/_base.py
+++ b/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(
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
index 17458cb76c16..ec44a89b68a4 100644
--- a/lib/matplotlib/cbook/__init__.py
+++ b/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):
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
index fe34034e4daf..86217ce89a88 100644
--- a/lib/matplotlib/tests/test_cbook.py
+++ b/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)
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
index ca25a9d069a6..fb0e6cc67596 100644
--- a/lib/mpl_toolkits/mplot3d/axes3d.py
+++ b/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)