Add snap to SpanSelector

Author: ericpre

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()
 

lib/matplotlib/widgets.py

@@ -2238,6 +2238,10 @@ 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, default: None
+        Snap the extents of the selector to the values defined in
+        ``snap_values``.
+
     Examples
     --------
     >>> import matplotlib.pyplot as plt
@@ -2259,7 +2263,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 +2282,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 +2582,16 @@ 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)."""
+        indices = np.empty_like(values, dtype="uint")
+        # take into account machine precision
+        eps = np.min(np.abs(np.diff(snap_values))) * 1e-12
+        for i, v in enumerate(values):
+            indices[i] = np.abs(snap_values - v + np.sign(v) * eps).argmin()
+        return snap_values[indices]
+
     @property
     def extents(self):
         """Return extents of the span selector."""
@@ -2591,6 +2606,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
@@ -2857,6 +2874,7 @@ def onselect(eclick: MouseEvent, erelease: MouseEvent)
     use_data_coordinates : bool, default: False
         If `True`, the "square" shape of the selector is defined in
         data coordinates instead of display coordinates.
+
     """