Merge pull request #21374 from ericpre/snap_selectors

oscargus · web-flow · commit a7b7260bf06c · 2022-04-21T11:31:07.000+02:00
Snap selectors
diff --git a/doc/users/next_whats_new/snap_selector.rst b/doc/users/next_whats_new/snap_selector.rst
@@ -0,0 +1,4 @@
+SpanSelector widget can now be snapped to specified values
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The SpanSelector widget can now be snapped to values specified by the *snap_values*
+argument.
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -67,8 +67,7 @@ def test_rectangle_selector():
 @pytest.mark.parametrize('spancoords', ['data', 'pixels'])
 @pytest.mark.parametrize('minspanx, x1', [[0, 10], [1, 10.5], [1, 11]])
 @pytest.mark.parametrize('minspany, y1', [[0, 10], [1, 10.5], [1, 11]])
-def test_rectangle_minspan(spancoords, minspanx, x1, minspany, y1):
-    ax = get_ax()
+def test_rectangle_minspan(ax, spancoords, minspanx, x1, minspany, y1):
     # attribute to track number of onselect calls
     ax._n_onselect = 0
 
@@ -924,6 +923,37 @@ def mean(vmin, vmax):
     assert ln2.stale is False
 
 
+def test_snapping_values_span_selector(ax):
+    def onselect(*args):
+        pass
+
+    tool = widgets.SpanSelector(ax, onselect, direction='horizontal',)
+    snap_function = tool._snap
+
+    snap_values = np.linspace(0, 5, 11)
+    values = np.array([-0.1, 0.1, 0.2, 0.5, 0.6, 0.7, 0.9, 4.76, 5.0, 5.5])
+    expect = np.array([00.0, 0.0, 0.0, 0.5, 0.5, 0.5, 1.0, 5.00, 5.0, 5.0])
+    values = snap_function(values, snap_values)
+    assert_allclose(values, expect)
+
+
+def test_span_selector_snap(ax):
+    def onselect(vmin, vmax):
+        ax._got_onselect = True
+
+    snap_values = np.arange(50) * 4
+
+    tool = widgets.SpanSelector(ax, onselect, direction='horizontal',
+                                snap_values=snap_values)
+    tool.extents = (17, 35)
+    assert tool.extents == (16, 36)
+
+    tool.snap_values = None
+    assert tool.snap_values is None
+    tool.extents = (17, 35)
+    assert tool.extents == (17, 35)
+
+
 def check_lasso_selector(**kwargs):
     ax = get_ax()
 
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -2238,6 +2238,9 @@ def on_select(min: float, max: float) -> Any
         If `True`, the event triggered outside the span selector will be
         ignored.
 
+    snap_values : 1D array-like, optional
+        Snap the selector edges to the given values.
+
     Examples
     --------
     >>> import matplotlib.pyplot as plt
@@ -2259,7 +2262,7 @@ def __init__(self, ax, onselect, direction, minspan=0, useblit=False,
                  props=None, onmove_callback=None, interactive=False,
                  button=None, handle_props=None, grab_range=10,
                  state_modifier_keys=None, drag_from_anywhere=False,
-                 ignore_event_outside=False):
+                 ignore_event_outside=False, snap_values=None):
 
         if state_modifier_keys is None:
             state_modifier_keys = dict(clear='escape',
@@ -2278,6 +2281,7 @@ def __init__(self, ax, onselect, direction, minspan=0, useblit=False,
 
         self.visible = True
         self._extents_on_press = None
+        self.snap_values = snap_values
 
         # self._pressv is deprecated and we don't use it internally anymore
         # but we maintain it until it is removed
@@ -2577,6 +2581,15 @@ def _contains(self, event):
         """Return True if event is within the patch."""
         return self._selection_artist.contains(event, radius=0)[0]
 
+    @staticmethod
+    def _snap(values, snap_values):
+        """Snap values to a given array values (snap_values)."""
+        # take into account machine precision
+        eps = np.min(np.abs(np.diff(snap_values))) * 1e-12
+        return tuple(
+            snap_values[np.abs(snap_values - v + np.sign(v) * eps).argmin()]
+            for v in values)
+
     @property
     def extents(self):
         """Return extents of the span selector."""
@@ -2591,6 +2604,8 @@ def extents(self):
     @extents.setter
     def extents(self, extents):
         # Update displayed shape
+        if self.snap_values is not None:
+            extents = tuple(self._snap(extents, self.snap_values))
         self._draw_shape(*extents)
         if self._interactive:
             # Update displayed handles
