Merge pull request #3937 from blink1073/rectangle_selector_upgrade

tacaswell · tacaswell · commit 8150af7306b8 · 2015-09-07T20:49:43.000-04:00
ENH: Rectangle Selector Upgrade
diff --git a/examples/widgets/rectangle_selector.py b/examples/widgets/rectangle_selector.py
@@ -46,6 +46,7 @@ def toggle_selector(event):
                                        drawtype='box', useblit=True,
                                        button=[1, 3],  # don't use middle button
                                        minspanx=5, minspany=5,
-                                       spancoords='pixels')
+                                       spancoords='pixels',
+                                       interactive=True)
 plt.connect('key_press_event', toggle_selector)
 plt.show()
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -11,8 +11,18 @@
 import matplotlib.pyplot as plt
 from matplotlib.testing.decorators import cleanup
 
+from numpy.testing import assert_allclose
 
-def get_event(ax, button=1, xdata=0, ydata=0, key=None, step=1):
+
+def get_ax():
+    fig, ax = plt.subplots(1, 1)
+    ax.plot([0, 200], [0, 200])
+    ax.set_aspect(1.0)
+    ax.figure.canvas.draw()
+    return ax
+
+
+def do_event(tool, etype, button=1, xdata=0, ydata=0, key=None, step=1):
     """
      *name*
         the event name
@@ -51,6 +61,7 @@ def get_event(ax, button=1, xdata=0, ydata=0, key=None, step=1):
     """
     event = mock.Mock()
     event.button = button
+    ax = tool.ax
     event.x, event.y = ax.transData.transform([(xdata, ydata),
                                                (xdata, ydata)])[00]
     event.xdata, event.ydata = xdata, ydata
@@ -60,32 +71,33 @@ def get_event(ax, button=1, xdata=0, ydata=0, key=None, step=1):
     event.step = step
     event.guiEvent = None
     event.name = 'Custom'
-    return event
+
+    func = getattr(tool, etype)
+    func(event)
 
 
 @cleanup
 def check_rectangle(**kwargs):
-    fig, ax = plt.subplots(1, 1)
-    ax.plot([0, 200], [0, 200])
-    ax.figure.canvas.draw()
+    ax = get_ax()
 
     def onselect(epress, erelease):
         ax._got_onselect = True
         assert epress.xdata == 100
         assert epress.ydata == 100
-        assert erelease.xdata == 200
-        assert erelease.ydata == 200
+        assert erelease.xdata == 199
+        assert erelease.ydata == 199
 
     tool = widgets.RectangleSelector(ax, onselect, **kwargs)
-    event = get_event(ax, xdata=100, ydata=100, button=1)
-    tool.press(event)
-
-    event = get_event(ax, xdata=125, ydata=125, button=1)
-    tool.onmove(event)
+    do_event(tool, 'press', xdata=100, ydata=100, button=1)
+    do_event(tool, 'onmove', xdata=199, ydata=199, button=1)
 
     # purposely drag outside of axis for release
-    event = get_event(ax, xdata=250, ydata=250, button=1)
-    tool.release(event)
+    do_event(tool, 'release', xdata=250, ydata=250, button=1)
+
+    if kwargs.get('drawtype', None) not in ['line', 'none']:
+        assert_allclose(tool.geometry,
+            [[100., 100, 199, 199, 100], [100, 199, 199, 100, 100]],
+            err_msg=tool.geometry)
 
     assert ax._got_onselect
 
@@ -99,11 +111,101 @@ def test_rectangle_selector():
     check_rectangle(rectprops=dict(fill=True))
 
 
+@cleanup
+def test_ellipse():
+    """For ellipse, test out the key modifiers"""
+    ax = get_ax()
+
+    def onselect(epress, erelease):
+        pass
+
+    tool = widgets.EllipseSelector(ax, onselect=onselect,
+                                   maxdist=10, interactive=True)
+    tool.extents = (100, 150, 100, 150)
+
+    # drag the rectangle
+    do_event(tool, 'press', xdata=10, ydata=10, button=1,
+                    key=' ')
+    do_event(tool, 'onmove', xdata=30, ydata=30, button=1)
+    do_event(tool, 'release', xdata=30, ydata=30, button=1)
+    assert tool.extents == (120, 170, 120, 170), tool.extents
+
+    # create from center
+    do_event(tool, 'on_key_press', xdata=100, ydata=100, button=1,
+                    key='control')
+    do_event(tool, 'press', xdata=100, ydata=100, button=1)
+    do_event(tool, 'onmove', xdata=125, ydata=125, button=1)
+    do_event(tool, 'release', xdata=125, ydata=125, button=1)
+    do_event(tool, 'on_key_release', xdata=100, ydata=100, button=1,
+                    key='control')
+    assert tool.extents == (75, 125, 75, 125), tool.extents
+
+    # create a square
+    do_event(tool, 'on_key_press', xdata=10, ydata=10, button=1,
+                    key='shift')
+    do_event(tool, 'press', xdata=10, ydata=10, button=1)
+    do_event(tool, 'onmove', xdata=35, ydata=30, button=1)
+    do_event(tool, 'release', xdata=35, ydata=30, button=1)
+    do_event(tool, 'on_key_release', xdata=10, ydata=10, button=1,
+                    key='shift')
+    extents = [int(e) for e in tool.extents]
+    assert extents == [10, 35, 10, 34]
+
+    # create a square from center
+    do_event(tool, 'on_key_press', xdata=100, ydata=100, button=1,
+                      key='ctrl+shift')
+    do_event(tool, 'press', xdata=100, ydata=100, button=1)
+    do_event(tool, 'onmove', xdata=125, ydata=130, button=1)
+    do_event(tool, 'release', xdata=125, ydata=130, button=1)
+    do_event(tool, 'on_key_release', xdata=100, ydata=100, button=1,
+                      key='ctrl+shift')
+    extents = [int(e) for e in tool.extents]
+    assert extents == [70, 129, 70, 130], extents
+
+    assert tool.geometry.shape == (2, 74)
+    assert_allclose(tool.geometry[:, 0], [70., 100])
+
+
+@cleanup
+def test_rectangle_handles():
+    ax = get_ax()
+
+    def onselect(epress, erelease):
+        pass
+
+    tool = widgets.RectangleSelector(ax, onselect=onselect,
+                                     maxdist=10, interactive=True)
+    tool.extents = (100, 150, 100, 150)
+
+    assert tool.corners == (
+        (100, 150, 150, 100), (100, 100, 150, 150))
+    assert tool.extents == (100, 150, 100, 150)
+    assert tool.edge_centers == (
+        (100, 125.0, 150, 125.0), (125.0, 100, 125.0, 150))
+    assert tool.extents == (100, 150, 100, 150)
+
+    # grab a corner and move it
+    do_event(tool, 'press', xdata=100, ydata=100)
+    do_event(tool, 'onmove', xdata=120, ydata=120)
+    do_event(tool, 'release', xdata=120, ydata=120)
+    assert tool.extents == (120, 150, 120, 150)
+
+    # grab the center and move it
+    do_event(tool, 'press', xdata=132, ydata=132)
+    do_event(tool, 'onmove', xdata=120, ydata=120)
+    do_event(tool, 'release', xdata=120, ydata=120)
+    assert tool.extents == (108, 138, 108, 138)
+
+    # create a new rectangle
+    do_event(tool, 'press', xdata=10, ydata=10)
+    do_event(tool, 'onmove', xdata=100, ydata=100)
+    do_event(tool, 'release', xdata=100, ydata=100)
+    assert tool.extents == (10, 100, 10, 100)
+
+
 @cleanup
 def check_span(*args, **kwargs):
-    fig, ax = plt.subplots(1, 1)
-    ax.plot([0, 200], [0, 200])
-    ax.figure.canvas.draw()
+    ax = get_ax()
 
     def onselect(vmin, vmax):
         ax._got_onselect = True
@@ -119,14 +221,9 @@ def onmove(vmin, vmax):
         kwargs['onmove_callback'] = onmove
 
     tool = widgets.SpanSelector(ax, onselect, *args, **kwargs)
-    event = get_event(ax, xdata=100, ydata=100, button=1)
-    tool.press(event)
-
-    event = get_event(ax, xdata=125, ydata=125, button=1)
-    tool.onmove(event)
-
-    event = get_event(ax, xdata=150, ydata=150, button=1)
-    tool.release(event)
+    do_event(tool, 'press', xdata=100, ydata=100, button=1)
+    do_event(tool, 'onmove', xdata=125, ydata=125, button=1)
+    do_event(tool, 'release', xdata=150, ydata=150, button=1)
 
     assert ax._got_onselect
 
@@ -142,24 +239,16 @@ def test_span_selector():
 
 @cleanup
 def check_lasso_selector(**kwargs):
-    fig, ax = plt.subplots(1, 1)
-    ax = plt.gca()
-    ax.plot([0, 200], [0, 200])
-    ax.figure.canvas.draw()
+    ax = get_ax()
 
     def onselect(verts):
         ax._got_onselect = True
         assert verts == [(100, 100), (125, 125), (150, 150)]
 
     tool = widgets.LassoSelector(ax, onselect, **kwargs)
-    event = get_event(ax, xdata=100, ydata=100, button=1)
-    tool.press(event)
-
-    event = get_event(ax, xdata=125, ydata=125, button=1)
-    tool.onmove(event)
-
-    event = get_event(ax, xdata=150, ydata=150, button=1)
-    tool.release(event)
+    do_event(tool, 'press', xdata=100, ydata=100, button=1)
+    do_event(tool, 'onmove', xdata=125, ydata=125, button=1)
+    do_event(tool, 'release', xdata=150, ydata=150, button=1)
 
     assert ax._got_onselect
 
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
