Merge pull request #19660 from dstansby/remove-poly-points

jklymak · web-flow · commit 93d1434d5646 · 2021-05-03T12:37:38.000-07:00
Allow PolygonSelector points to be removed
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 completing 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,57 @@ 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)
+
+
+def test_polygon_selector_redraw():
+    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 completed, now remove first two verts
+                      polygon_remove_vertex(*verts[1]) +
+                      polygon_remove_vertex(*verts[2]) +
+                      # At this point the tool should be reset so we can add
+                      # more vertices
+                      polygon_place_vertex(*verts[1]))
+
+    ax = get_ax()
+
+    def onselect(vertices):
+        pass
+
+    tool = widgets.PolygonSelector(ax, onselect)
+    for (etype, event_args) in event_sequence:
+        do_event(tool, etype, **event_args)
+    # After removing two verts, only one remains, and the
+    # selector should be automatically resete
+    assert tool.verts == verts[0: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 individual vertices
+    can be moved by clicking and dragging with the left mouse button, or
+    removed 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 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.
 
     For the selector to remain responsive you must keep a reference to it.
 
@@ -2789,6 +2795,12 @@ class PolygonSelector(_SelectorWidget):
     Examples
     --------
     :doc:`/gallery/widgets/polygon_selector_demo`
+
+    Notes
+    -----
+    If only one point remains after removing points, the selector reverts to an
+    incomplete state and you can start drawing a new polygon from the existing
+    point.
     """
 
     def __init__(self, ax, onselect, useblit=False,
@@ -2827,6 +2839,33 @@ 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."""
+        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 incomplete 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 +2882,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.
