Add set_props and set_handle_props to selectors.

ericpre · ericpre · commit 15beb47d9585 · 2021-07-27T11:07:30.000+01:00
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -90,6 +90,36 @@ def onselect(epress, erelease):
     assert tool.center == (180, 190)
 
 
+def test_rectangle_selector_set_props_handle_props():
+    ax = get_ax()
+
+    def onselect(epress, erelease):
+        pass
+
+    tool = widgets.RectangleSelector(ax, onselect, interactive=True,
+                                     props=dict(facecolor='b', alpha=0.2),
+                                     handle_props=dict(alpha=0.5))
+    # Create rectangle
+    do_event(tool, 'press', xdata=0, ydata=10, button=1)
+    do_event(tool, 'onmove', xdata=100, ydata=120, button=1)
+    do_event(tool, 'release', xdata=100, ydata=120, button=1)
+
+    artist = tool.artists[0]
+    assert artist.get_facecolor() == mcolors.to_rgba('b', alpha=0.2)
+    props = dict(facecolor='r', alpha=0.3)
+    tool.set_props(**props)
+    assert artist.get_facecolor() == mcolors.to_rgba(*props.values())
+
+    for artist in tool._handles_artists:
+        assert artist.get_markeredgecolor() == 'black'
+        assert artist.get_alpha() == 0.5
+    handle_props = dict(markeredgecolor='r', alpha=0.3)
+    tool.set_handle_props(**handle_props)
+    for artist in tool._handles_artists:
+        assert artist.get_markeredgecolor() == 'r'
+        assert artist.get_alpha() == 0.3
+
+
 def test_ellipse():
     """For ellipse, test out the key modifiers"""
     ax = get_ax()
@@ -185,9 +215,9 @@ def onselect(epress, erelease):
 
     # Check that marker_props worked.
     assert mcolors.same_color(
-        tool._corner_handles.artist.get_markerfacecolor(), 'r')
+        tool._corner_handles.artists[0].get_markerfacecolor(), 'r')
     assert mcolors.same_color(
-        tool._corner_handles.artist.get_markeredgecolor(), 'b')
+        tool._corner_handles.artists[0].get_markeredgecolor(), 'b')
 
 
 def check_span(*args, **kwargs):
@@ -281,6 +311,36 @@ def onselect(epress, erelease):
         tool.direction = 'invalid_string'
 
 
+def test_span_selector_set_props_handle_props():
+    ax = get_ax()
+
+    def onselect(epress, erelease):
+        pass
+
+    tool = widgets.SpanSelector(ax, onselect, 'horizontal', interactive=True,
+                                props=dict(facecolor='b', alpha=0.2),
+                                handle_props=dict(alpha=0.5))
+    # Create rectangle
+    do_event(tool, 'press', xdata=0, ydata=10, button=1)
+    do_event(tool, 'onmove', xdata=100, ydata=120, button=1)
+    do_event(tool, 'release', xdata=100, ydata=120, button=1)
+
+    artist = tool.artists[0]
+    assert artist.get_facecolor() == mcolors.to_rgba('b', alpha=0.2)
+    props = dict(facecolor='r', alpha=0.3)
+    tool.set_props(**props)
+    assert artist.get_facecolor() == mcolors.to_rgba(*props.values())
+
+    for artist in tool._handles_artists:
+        assert artist.get_color() == 'b'
+        assert artist.get_alpha() == 0.5
+    handle_props = dict(color='r', alpha=0.3)
+    tool.set_handle_props(**handle_props)
+    for artist in tool._handles_artists:
+        assert artist.get_color() == 'r'
+        assert artist.get_alpha() == 0.3
+
+
 def test_tool_line_handle():
     ax = get_ax()
 
@@ -658,6 +718,45 @@ def test_polygon_selector():
     check_polygon_selector(event_sequence, expected_result, 1)
 
 
+def test_polygon_selector_set_props_handle_props():
+    ax = get_ax()
+
+    ax._selections_count = 0
+
+    def onselect(vertices):
+        ax._selections_count += 1
+        ax._current_result = vertices
+
+    tool = widgets.PolygonSelector(ax, onselect,
+                                   props=dict(color='b', alpha=0.2),
+                                   handle_props=dict(alpha=0.5))
+
+    event_sequence = (polygon_place_vertex(50, 50)
+                      + polygon_place_vertex(150, 50)
+                      + polygon_place_vertex(50, 150)
+                      + polygon_place_vertex(50, 50))
+
+    for (etype, event_args) in event_sequence:
+        do_event(tool, etype, **event_args)
+
+    artist = tool.artists[0]
+    assert artist.get_color() == 'b'
+    assert artist.get_alpha() == 0.2
+    props = dict(color='r', alpha=0.3)
+    tool.set_props(**props)
+    assert artist.get_color() == props['color']
+    assert artist.get_alpha() == props['alpha']
+
+    for artist in tool._handles_artists:
+        assert artist.get_color() == 'b'
+        assert artist.get_alpha() == 0.5
+    handle_props = dict(color='r', alpha=0.3)
+    tool.set_handle_props(**handle_props)
+    for artist in tool._handles_artists:
+        assert artist.get_color() == 'r'
+        assert artist.get_alpha() == 0.3
+
+
 @pytest.mark.parametrize(
     "horizOn, vertOn",
     [(True, True), (True, False), (False, True)],
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -2013,8 +2013,31 @@ def add_artist(self, artist):
 
     @property
     def artists(self):
+        """Tuple of the artists of the selector"""
         return tuple(self._artists)
 
+    def set_props(self, **props):
+        """
+        Set the properties of the selector artist. See the `props` argument
+        in the selector docstring to know which properties are supported.
+        """
+        self.artists[0].set(**props)
+        self.update()
+        self._props = props
+
+    def set_handle_props(self, **handle_props):
+        """
+        Set the properties of the handles selector artist. See the
+        `handle_props` argument in the selector docstring to know which
+        properties are supported.
+        """
+        if not hasattr(self, '_handles_artists'):
+            raise NotImplementedError("This selector doesn't have handles.")
+        for handle in self._handles_artists:
+            handle.set(**handle_props)
+        self.update()
+        self._handle_props = handle_props
+
 
 class SpanSelector(_SelectorWidget):
     """
@@ -2137,7 +2160,7 @@ def __init__(self, ax, onselect, direction, minspan=0, useblit=False,
 
         if self._interactive:
             self._edge_order = ['min', 'max']
-            self._setup_edge_handle(handle_props)
+            self._setup_edge_handles(handle_props)
 
         self._active_handle = None
 
@@ -2187,7 +2210,7 @@ def new_axes(self, ax):
         else:
             self.add_artist(rect_artist)
 
-    def _setup_edge_handle(self, props):
+    def _setup_edge_handles(self, props):
         # Define initial position using the axis bounds to keep the same bounds
         if self.direction == 'horizontal':
             positions = self.ax.get_xbound()
@@ -2267,7 +2290,7 @@ def direction(self, direction):
             self._direction = direction
             self.new_axes(self.ax)
             if self._interactive:
-                self._setup_edge_handle(self._edge_handles._line_props)
+                self._setup_edge_handles(self._edge_handles._line_props)
         else:
             self._direction = direction
 
@@ -2409,6 +2432,10 @@ def extents(self, extents):
         self.set_visible(self.visible)
         self.update()
 
+    @property
+    def _handles_artists(self):
+        return self._edge_handles.artists
+
 
 class ToolLineHandles:
     """
@@ -2545,7 +2572,6 @@ def __init__(self, ax, x, y, marker='o', marker_props=None, useblit=True):
                  **cbook.normalize_kwargs(marker_props, Line2D._alias_map)}
         self._markers = Line2D(x, y, animated=useblit, **props)
         self.ax.add_line(self._markers)
-        self.artist = self._markers
 
     @property
     def x(self):
@@ -2555,6 +2581,10 @@ def x(self):
     def y(self):
         return self._markers.get_ydata()
 
+    @property
+    def artists(self):
+        return (self._markers, )
+
     def set_data(self, pts, y=None):
         """Set x and y positions of handles."""
         if y is not None:
@@ -2765,9 +2795,7 @@ def __init__(self, ax, onselect, drawtype='box',
 
             self._active_handle = None
 
-            self._artists.extend([self._center_handle.artist,
-                                  self._corner_handles.artist,
-                                  self._edge_handles.artist])
+            self._artists.extend(self._handles_artists)
 
         self._extents_on_press = None
 
@@ -2783,6 +2811,11 @@ def __init__(self, ax, onselect, drawtype='box',
         property(lambda self: self.grab_range,
                  lambda self, value: setattr(self, "grab_range", value)))
 
+    @property
+    def _handles_artists(self):
+        return (self._center_handle.artists + self._corner_handles.artists +
+                self._edge_handles.artists)
+
     def _press(self, event):
         """Button press event handler."""
         # make the drawn box/line visible get the click-coordinates,
@@ -3302,7 +3335,7 @@ def __init__(self, ax, onselect, useblit=False,
         self._active_handle_idx = -1
         self.grab_range = grab_range
 
-        self._artists = [line, self._polygon_handles.artist]
+        self._artists = [line] + list(self._handles_artists)
         self.set_visible(True)
 
     vertex_select_radius = _api.deprecated("3.5", name="vertex_select_radius",
@@ -3315,6 +3348,10 @@ def __init__(self, ax, onselect, useblit=False,
     def _nverts(self):
         return len(self._xs)
 
+    @property
+    def _handles_artists(self):
+        return self._polygon_handles.artists
+
     def _remove_vertex(self, i):
         """Remove vertex with index i."""
         if (self._nverts > 2 and
