Small cleanup to VertexSelector.

anntzer · anntzer · commit a5998a560c1b · 2022-07-08T15:19:10.000+02:00
Make the class picklable, by making `canvas` a property and using a
picklable callback connect.  Also some docstring fixes.
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -1515,16 +1515,15 @@ class VertexSelector:
     Derived classes should override the `process_selected` method to do
     something with the picks.
 
-    Here is an example which highlights the selected verts with red
-    circles::
+    Here is an example which highlights the selected verts with red circles::
 
         import numpy as np
         import matplotlib.pyplot as plt
         import matplotlib.lines as lines
 
         class HighlightSelected(lines.VertexSelector):
             def __init__(self, line, fmt='ro', **kwargs):
-                lines.VertexSelector.__init__(self, line)
+                super().__init__(line)
                 self.markers, = self.axes.plot([], [], fmt, **kwargs)
 
             def process_selected(self, ind, xs, ys):
@@ -1537,27 +1536,29 @@ def process_selected(self, ind, xs, ys):
 
         selector = HighlightSelected(line)
         plt.show()
-
     """
+
     def __init__(self, line):
         """
-        Initialize the class with a `.Line2D`.  The line should already be
-        added to an `~.axes.Axes` and should have the picker property set.
+        Parameters
+        ----------
+        line : `.Line2D`
+            The line must already have been added to an `~.axes.Axes` and must
+            have its picker property set.
         """
         if line.axes is None:
             raise RuntimeError('You must first add the line to the Axes')
-
         if line.get_picker() is None:
             raise RuntimeError('You must first set the picker property '
                                'of the line')
-
         self.axes = line.axes
         self.line = line
-        self.canvas = self.axes.figure.canvas
-        self.cid = self.canvas.mpl_connect('pick_event', self.onpick)
-
+        self.cid = self.canvas.callbacks._connect_picklable(
+            'pick_event', self.onpick)
         self.ind = set()
 
+    canvas = property(lambda self: self.axes.figure.canvas)
+
     def process_selected(self, ind, xs, ys):
         """
         Default "do nothing" implementation of the `process_selected` method.
diff --git a/lib/matplotlib/tests/test_pickle.py b/lib/matplotlib/tests/test_pickle.py
@@ -8,6 +8,7 @@
 from matplotlib import cm
 from matplotlib.testing.decorators import check_figures_equal
 from matplotlib.dates import rrulewrapper
+from matplotlib.lines import VertexSelector
 import matplotlib.pyplot as plt
 import matplotlib.transforms as mtransforms
 import matplotlib.figure as mfigure
@@ -231,3 +232,8 @@ def test_dynamic_norm():
         mpl.scale.LogitScale, mpl.colors.Normalize)()
     assert type(pickle.loads(pickle.dumps(logit_norm_instance))) \
         == type(logit_norm_instance)
+
+
+def test_vertexselector():
+    line, = plt.plot([0, 1], picker=True)
+    pickle.loads(pickle.dumps(VertexSelector(line)))
