Allow PolygonSelector points to be removed

dstansby · dstansby · commit 070b9129ddd6 · 2021-05-03T10:57:04.000+01:00
diff --git a/doc/users/next_whats_new/polygons_selector_remove_points.rst b/doc/users/next_whats_new/polygons_selector_remove_points.rst
@@ -0,0 +1,4 @@
+Removing points on a PolygonSelector
+------------------------------------
+After fishing drawing a `~matplotlib.widgets.PolygonSelector`, individual
+points can now be removed by right-clicking on them.
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -416,6 +416,12 @@ def polygon_place_vertex(xdata, ydata):
             ('release', dict(xdata=xdata, ydata=ydata))]
 
 
+def polygon_remove_vertex(xdata, ydata):
+    return [('onmove', dict(xdata=xdata, ydata=ydata)),
+            ('press', dict(xdata=xdata, ydata=ydata, button=3)),
+            ('release', dict(xdata=xdata, ydata=ydata, button=3))]
+
+
 def test_polygon_selector():
     # Simple polygon
     expected_result = [(50, 50), (150, 50), (50, 150)]
@@ -564,3 +570,31 @@ def onselect(verts):
     tool = widgets.RectangleSelector(ax_test, onselect,
                                      rectprops={'visible': False})
     tool.extents = (0.2, 0.8, 0.3, 0.7)
+
+
+# Change the order that the extra point is inserted in
+@pytest.mark.parametrize('idx', [1, 2, 3])
+def test_polygon_selector_remove(idx):
+    verts = [(50, 50), (150, 50), (50, 150)]
+    event_sequence = [polygon_place_vertex(*verts[0]),
+                      polygon_place_vertex(*verts[1]),
+                      polygon_place_vertex(*verts[2]),
+                      # Finish the polygon
+                      polygon_place_vertex(*verts[0])]
+    # Add an extra point
+    event_sequence.insert(idx, polygon_place_vertex(200, 200))
+    # Remove the extra point
+    event_sequence.append(polygon_remove_vertex(200, 200))
+    # Flatten list of lists
+    event_sequence = sum(event_sequence, [])
+    check_polygon_selector(event_sequence, verts, 2)
+
+
+def test_polygon_selector_remove_first_point():
+    verts = [(50, 50), (150, 50), (50, 150)]
+    event_sequence = (polygon_place_vertex(*verts[0]) +
+                      polygon_place_vertex(*verts[1]) +
+                      polygon_place_vertex(*verts[2]) +
+                      polygon_place_vertex(*verts[0]) +
+                      polygon_remove_vertex(*verts[0]))
+    check_polygon_selector(event_sequence, verts[1:], 2)
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -2759,11 +2759,17 @@ class PolygonSelector(_SelectorWidget):
     Select a polygon region of an axes.
 
     Place vertices with each mouse click, and make the selection by completing
-    the polygon (clicking on the first vertex). Hold the *ctrl* key and click
-    and drag a vertex to reposition it (the *ctrl* key is not necessary if the
-    polygon has already been completed). Hold the *shift* key and click and
-    drag anywhere in the axes to move all vertices. Press the *esc* key to
-    start a new polygon.
+    the polygon (clicking on the first vertex). Once drawn indiviual vertices
+    can be moved by clicking and dragging with the left mouse button, or
+    deleted by clicking the right mouse button.
+
+    In addition, the following modifier keys can be used:
+
+    - Hold *ctrl* and click and drag a vertex to reposition it before the
+      polygon has been copleted.
+    - Hold the *shift* key and click and drag anywhere in the axes to move
+      all vertices.
+    - Press the *esc* key to start a new polygon.
 
     For the selector to remain responsive you must keep a reference to it.
 
@@ -2827,6 +2833,34 @@ def __init__(self, ax, onselect, useblit=False,
         self.artists = [self.line, self._polygon_handles.artist]
         self.set_visible(True)
 
+    @property
+    def _nverts(self):
+        return len(self._xs)
+
+    def _remove_vertex(self, i):
+        """Remove vertex with index i."""
+        print(self._nverts)
+        if (self._nverts > 2 and
+                self._polygon_completed and
+                i in (0, self._nverts - 1)):
+            # If selecting the first or final vertex, remove both first and
+            # last vertex as they are the same for a closed polygon
+            self._xs.pop(0)
+            self._ys.pop(0)
+            self._xs.pop(-1)
+            self._ys.pop(-1)
+            # Close the polygon again by appending the new first vertex to the
+            # end
+            self._xs.append(self._xs[0])
+            self._ys.append(self._ys[0])
+        else:
+            self._xs.pop(i)
+            self._ys.pop(i)
+        if self._nverts <= 2:
+            # If only one point left, return to un-complete state to let user
+            # start drawing again
+            self._polygon_completed = False
+
     def _press(self, event):
         """Button press event handler."""
         # Check for selection of a tool handle.
@@ -2843,6 +2877,9 @@ def _release(self, event):
         """Button release event handler."""
         # Release active tool handle.
         if self._active_handle_idx >= 0:
+            if event.button == 3:
+                self._remove_vertex(self._active_handle_idx)
+                self._draw_polygon()
             self._active_handle_idx = -1
 
         # Complete the polygon.
