matplotlib · phobson · Apr 13, 2017 · Mar 28, 2017 · Mar 30, 2017 · Mar 30, 2017
diff --git a/examples/widgets/lasso_selector_demo.py b/examples/widgets/lasso_selector_demo.py
@@ -11,9 +11,10 @@
 class SelectFromCollection(object):
     """Select indices from a matplotlib collection using `LassoSelector`.
 
-    Selected indices are saved in the `ind` attribute. This tool highlights
-    selected points by fading them out (i.e., reducing their alpha values).
-    If your collection has alpha < 1, this tool will permanently alter them.
+    Selected indices are saved in the `ind` attribute. This tool fades out the
+    points that are not part of the selection (i.e., reduces their alpha
+    values). If your collection has alpha < 1, this tool will permanently
+    alter the alpha values.
 
     Note that this tool selects collection objects based on their *origins*
     (i.e., `offsets`).
@@ -44,14 +45,14 @@ def __init__(self, ax, collection, alpha_other=0.3):
         if len(self.fc) == 0:
             raise ValueError('Collection must have a facecolor')
         elif len(self.fc) == 1:
-            self.fc = np.tile(self.fc, self.Npts).reshape(self.Npts, -1)
+            self.fc = np.tile(self.fc, (self.Npts, 1))
 
         self.lasso = LassoSelector(ax, onselect=self.onselect)
         self.ind = []
 
     def onselect(self, verts):
         path = Path(verts)
-        self.ind = np.nonzero([path.contains_point(xy) for xy in self.xys])[0]
+        self.ind = np.nonzero(path.contains_points(self.xys))[0]
         self.fc[:, -1] = self.alpha_other
         self.fc[self.ind, -1] = 1
         self.collection.set_facecolors(self.fc)

diff --git a/examples/widgets/polygon_selector_demo.py b/examples/widgets/polygon_selector_demo.py
@@ -0,0 +1,87 @@
+
+import numpy as np
+
+from matplotlib.widgets import PolygonSelector
+from matplotlib.path import Path
+
+
+class SelectFromCollection(object):
+    """Select indices from a matplotlib collection using `PolygonSelector`.
+
+    Selected indices are saved in the `ind` attribute. This tool fades out the
+    points that are not part of the selection (i.e., reduces their alpha
+    values). If your collection has alpha < 1, this tool will permanently
+    alter the alpha values.
+
+    Note that this tool selects collection objects based on their *origins*
+    (i.e., `offsets`).
+
+    Parameters
+    ----------
+    ax : :class:`~matplotlib.axes.Axes`
+        Axes to interact with.
+
+    collection : :class:`matplotlib.collections.Collection` subclass
+        Collection you want to select from.
+
+    alpha_other : 0 <= float <= 1
+        To highlight a selection, this tool sets all selected points to an
+        alpha value of 1 and non-selected points to `alpha_other`.
+    """
+
+    def __init__(self, ax, collection, alpha_other=0.3):
+        self.canvas = ax.figure.canvas
+        self.collection = collection
+        self.alpha_other = alpha_other
+
+        self.xys = collection.get_offsets()
+        self.Npts = len(self.xys)
+
+        # Ensure that we have separate colors for each object
+        self.fc = collection.get_facecolors()
+        if len(self.fc) == 0:
+            raise ValueError('Collection must have a facecolor')
+        elif len(self.fc) == 1:
+            self.fc = np.tile(self.fc, (self.Npts, 1))
+
+        self.poly = PolygonSelector(ax, self.onselect)
+        self.ind = []
+
+    def onselect(self, verts):
+        path = Path(verts)
+        self.ind = np.nonzero(path.contains_points(self.xys))[0]
+        self.fc[:, -1] = self.alpha_other
+        self.fc[self.ind, -1] = 1
+        self.collection.set_facecolors(self.fc)
+        self.canvas.draw_idle()
+
+    def disconnect(self):
+        self.poly.disconnect_events()
+        self.fc[:, -1] = 1
+        self.collection.set_facecolors(self.fc)
+        self.canvas.draw_idle()
+
+
+if __name__ == '__main__':
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots()
+    grid_size = 5
+    grid_x = np.tile(np.arange(grid_size), grid_size)
+    grid_y = np.repeat(np.arange(grid_size), grid_size)
+    pts = ax.scatter(grid_x, grid_y)
+
+    selector = SelectFromCollection(ax, pts)
+
+    print("Select points in the figure by enclosing them within a polygon.")
+    print("Press the 'esc' key to start a new polygon.")
+    print("Try holding the 'shift' key to move all of the vertices.")
+    print("Try holding the 'ctrl' key to move a single vertex.")
+
+    plt.show()
+
+    selector.disconnect()
+
+    # After figure is closed print the coordinates of the selected points
+    print('\nSelected points:')
+    print(selector.xys[selector.ind])
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -313,3 +313,139 @@ def test_slider_valmin_valmax():
     slider = widgets.Slider(ax=ax, label='', valmin=0.0, valmax=24.0,
                             valinit=25.0)
     assert slider.val == slider.valmax
+
+
+def check_polygon_selector(event_sequence, expected_result, selections_count):
+    """Helper function to test Polygon Selector
+
+    Parameters
+    ----------
+    event_sequence : list of tuples (etype, dict())
+        A sequence of events to perform. The sequence is a list of tuples
+        where the first element of the tuple is an etype (e.g., 'onmove',
+        'press', etc.), and the second element of the tuple is a dictionary of
+         the arguments for the event (e.g., xdata=5, key='shift', etc.).
+    expected_result : list of vertices (xdata, ydata)
+        The list of vertices that are expected to result from the event
+        sequence.
+    selections_count : int
+        Wait for the tool to call its `onselect` function `selections_count`
+        times, before comparing the result to the `expected_result`
+    """
+    ax = get_ax()
+
+    ax._selections_count = 0
+
+    def onselect(vertices):
+        ax._selections_count += 1
+        ax._current_result = vertices
+
+    tool = widgets.PolygonSelector(ax, onselect)
+
+    for (etype, event_args) in event_sequence:
+        do_event(tool, etype, **event_args)
+
+    assert ax._selections_count == selections_count
+    assert ax._current_result == expected_result
+
+
+def polygon_place_vertex(xdata, ydata):
+    return [('onmove', dict(xdata=xdata, ydata=ydata)),
+            ('press', dict(xdata=xdata, ydata=ydata)),
+            ('release', dict(xdata=xdata, ydata=ydata))]
+
+
+def test_polygon_selector():
+    # Simple polygon
+    expected_result = [(50, 50), (150, 50), (50, 150)]
+    event_sequence = (polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(150, 50)
+                      + polygon_place_vertex(50, 150)
+                      + polygon_place_vertex(50, 50))
+    check_polygon_selector(event_sequence, expected_result, 1)
+
+    # Move first vertex before completing the polygon.
+    expected_result = [(75, 50), (150, 50), (50, 150)]
+    event_sequence = (polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(150, 50)
+                      + [('on_key_press', dict(key='control')),
+                         ('onmove', dict(xdata=50, ydata=50)),
+                         ('press', dict(xdata=50, ydata=50)),
+                         ('onmove', dict(xdata=75, ydata=50)),
+                         ('release', dict(xdata=75, ydata=50)),
+                         ('on_key_release', dict(key='control'))]
+                      + polygon_place_vertex(50, 150)
+                      + polygon_place_vertex(75, 50))
+    check_polygon_selector(event_sequence, expected_result, 1)
+
+    # Move first two vertices at once before completing the polygon.
+    expected_result = [(50, 75), (150, 75), (50, 150)]
+    event_sequence = (polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(150, 50)
+                      + [('on_key_press', dict(key='shift')),
+                         ('onmove', dict(xdata=100, ydata=100)),
+                         ('press', dict(xdata=100, ydata=100)),
+                         ('onmove', dict(xdata=100, ydata=125)),
+                         ('release', dict(xdata=100, ydata=125)),
+                         ('on_key_release', dict(key='shift'))]
+                      + polygon_place_vertex(50, 150)
+                      + polygon_place_vertex(50, 75))
+    check_polygon_selector(event_sequence, expected_result, 1)
+
+    # Move first vertex after completing the polygon.
+    expected_result = [(75, 50), (150, 50), (50, 150)]
+    event_sequence = (polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(150, 50)
+                      + polygon_place_vertex(50, 150)
+                      + polygon_place_vertex(50, 50)
+                      + [('onmove', dict(xdata=50, ydata=50)),
+                         ('press', dict(xdata=50, ydata=50)),
+                         ('onmove', dict(xdata=75, ydata=50)),
+                         ('release', dict(xdata=75, ydata=50))])
+    check_polygon_selector(event_sequence, expected_result, 2)
+
+    # Move all vertices after completing the polygon.
+    expected_result = [(75, 75), (175, 75), (75, 175)]
+    event_sequence = (polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(150, 50)
+                      + polygon_place_vertex(50, 150)
+                      + polygon_place_vertex(50, 50)
+                      + [('on_key_press', dict(key='shift')),
+                         ('onmove', dict(xdata=100, ydata=100)),
+                         ('press', dict(xdata=100, ydata=100)),
+                         ('onmove', dict(xdata=125, ydata=125)),
+                         ('release', dict(xdata=125, ydata=125)),
+                         ('on_key_release', dict(key='shift'))])
+    check_polygon_selector(event_sequence, expected_result, 2)
+
+    # Try to move a vertex and move all before placing any vertices.
+    expected_result = [(50, 50), (150, 50), (50, 150)]
+    event_sequence = ([('on_key_press', dict(key='control')),
+                       ('onmove', dict(xdata=100, ydata=100)),
+                       ('press', dict(xdata=100, ydata=100)),
+                       ('onmove', dict(xdata=125, ydata=125)),
+                       ('release', dict(xdata=125, ydata=125)),
+                       ('on_key_release', dict(key='control')),
+                       ('on_key_press', dict(key='shift')),
+                       ('onmove', dict(xdata=100, ydata=100)),
+                       ('press', dict(xdata=100, ydata=100)),
+                       ('onmove', dict(xdata=125, ydata=125)),
+                       ('release', dict(xdata=125, ydata=125)),
+                       ('on_key_release', dict(key='shift'))]
+                      + polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(150, 50)
+                      + polygon_place_vertex(50, 150)
+                      + polygon_place_vertex(50, 50))
+    check_polygon_selector(event_sequence, expected_result, 1)
+
+    # Try to place vertex out-of-bounds, then reset, and start a new polygon.
+    expected_result = [(50, 50), (150, 50), (50, 150)]
+    event_sequence = (polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(250, 50)
+                      + [('on_key_press', dict(key='escape')),
+                         ('on_key_release', dict(key='escape'))]
+                      + polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(150, 50)
+                      + polygon_place_vertex(50, 150)
+                      + polygon_place_vertex(50, 50))
+    check_polygon_selector(event_sequence, expected_result, 1)