Implement rotation selector

ericpre · ericpre · commit 9350a023726b · 2021-08-25T09:37:31.000+01:00
diff --git a/lib/matplotlib/patches.py b/lib/matplotlib/patches.py
@@ -709,7 +709,8 @@ def __str__(self):
         return fmt % pars
 
     @docstring.dedent_interpd
-    def __init__(self, xy, width, height, angle=0.0, **kwargs):
+    def __init__(self, xy, width, height, angle=0.0,
+                 rotate_around_center=False, **kwargs):
         """
         Parameters
         ----------
@@ -720,7 +721,12 @@ def __init__(self, xy, width, height, angle=0.0, **kwargs):
         height : float
             Rectangle height.
         angle : float, default: 0
-            Rotation in degrees anti-clockwise about *xy*.
+            Rotation in degrees anti-clockwise about *xy* if
+            *rotate_around_center* if False, otherwise rotate around the
+            center of the rectangle
+        rotate_around_center : bool, default: False
+            If True, the rotation is performed around the center of the
+            rectangle.
 
         Other Parameters
         ----------------
@@ -733,6 +739,7 @@ def __init__(self, xy, width, height, angle=0.0, **kwargs):
         self._width = width
         self._height = height
         self.angle = float(angle)
+        self.rotate_around_center = rotate_around_center
         self._convert_units()  # Validate the inputs.
 
     def get_path(self):
@@ -753,9 +760,14 @@ def get_patch_transform(self):
         # important to call the accessor method and not directly access the
         # transformation member variable.
         bbox = self.get_bbox()
+        if self.rotate_around_center:
+            width, height = bbox.x1 - bbox.x0, bbox.y1 - bbox.y0
+            rotation_point = bbox.x0 + width / 2., bbox.y0 + height / 2.
+        else:
+            rotation_point = bbox.x0, bbox.y0
         return (transforms.BboxTransformTo(bbox)
                 + transforms.Affine2D().rotate_deg_around(
-                    bbox.x0, bbox.y0, self.angle))
+                    *rotation_point, self.angle))
 
     def get_x(self):
         """Return the left coordinate of the rectangle."""
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -510,11 +510,10 @@ def onselect(epress, erelease):
                                                    'markeredgecolor': 'b'})
     tool.extents = (100, 150, 100, 150)
 
-    assert tool.corners == (
-        (100, 150, 150, 100), (100, 100, 150, 150))
+    assert_allclose(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_allclose(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
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -20,6 +20,7 @@
 from . import _api, backend_tools, cbook, colors, ticker
 from .lines import Line2D
 from .patches import Circle, Rectangle, Ellipse
+from .transforms import Affine2D
 
 
 class LockDraw:
@@ -1797,7 +1798,8 @@ def __init__(self, ax, onselect, useblit=False, button=None,
 
         self._state_modifier_keys = dict(move=' ', clear='escape',
                                          square='shift', center='control',
-                                         data_coordinates='d')
+                                         data_coordinates='d',
+                                         rotate='r')
         self._state_modifier_keys.update(state_modifier_keys or {})
 
         self.background = None
@@ -1932,8 +1934,9 @@ def press(self, event):
             key = event.key or ''
             key = key.replace('ctrl', 'control')
             # move state is locked in on a button press
-            if key == self._state_modifier_keys['move']:
-                self._state.add('move')
+            for action in ['move']:
+                if key == self._state_modifier_keys[action]:
+                    self._state.add(action)
             self._press(event)
             return True
         return False
@@ -1985,14 +1988,15 @@ def on_key_press(self, event):
                     artist.set_visible(False)
                 self.update()
                 return
-            if key == 'd' and key in self._state_modifier_keys.values():
-                modifier = 'data_coordinates'
-                if modifier in self._default_state:
-                    self._default_state.remove(modifier)
-                else:
-                    self.add_default_state(modifier)
+            for state in ['rotate', 'data_coordinates']:
+                if key == self._state_modifier_keys[state]:
+                    if state in self._default_state:
+                        self._default_state.remove(state)
+                    else:
+                        self.add_default_state(state)
             for (state, modifier) in self._state_modifier_keys.items():
-                if modifier in key:
+                # Multiple keys are string concatenated using '+'
+                if modifier in key.split('+'):
                     self._state.add(state)
             self._on_key_press(event)
 
@@ -2193,7 +2197,8 @@ def __init__(self, ax, onselect, direction, minspan=0, useblit=False,
             state_modifier_keys = dict(clear='escape',
                                        square='not-applicable',
                                        center='not-applicable',
-                                       data_coordinates='not-applicable')
+                                       data_coordinates='not-applicable',
+                                       rotate='not-applicable')
         super().__init__(ax, onselect, useblit=useblit, button=button,
                          state_modifier_keys=state_modifier_keys)
 
@@ -2764,6 +2769,7 @@ def onselect(eclick: MouseEvent, erelease: MouseEvent)
         - "center": change the shape around its center, default: "ctrl".
         - "data_coordinates": define if data or figure coordinates should be
           used to define the square shape, default: "d"
+        - "rotate": Rotate the shape around its center, default: "r".
 
         "square" and "center" can be combined. The square shape can be defined
         in data or figure coordinates as determined by the ``data_coordinates``
@@ -2811,8 +2817,6 @@ class RectangleSelector(_SelectorWidget):
     See also: :doc:`/gallery/widgets/rectangle_selector`
     """
 
-    _shape_klass = Rectangle
-
     @_api.rename_parameter("3.5", "maxdist", "grab_range")
     @_api.rename_parameter("3.5", "marker_props", "handle_props")
     @_api.rename_parameter("3.5", "rectprops", "props")
@@ -2831,6 +2835,7 @@ def __init__(self, ax, onselect, drawtype='box',
         self._interactive = interactive
         self.drag_from_anywhere = drag_from_anywhere
         self.ignore_event_outside = ignore_event_outside
+        self._rotation = 0
 
         if drawtype == 'none':  # draw a line but make it invisible
             _api.warn_deprecated(
@@ -2848,8 +2853,7 @@ def __init__(self, ax, onselect, drawtype='box',
             props['animated'] = self.useblit
             self.visible = props.pop('visible', self.visible)
             self._props = props
-            to_draw = self._shape_klass((0, 0), 0, 1, visible=False,
-                                        **self._props)
+            to_draw = self._init_shape(**self._props)
             self.ax.add_patch(to_draw)
         if drawtype == 'line':
             _api.warn_deprecated(
@@ -2921,6 +2925,10 @@ def _handles_artists(self):
         return (*self._center_handle.artists, *self._corner_handles.artists,
                 *self._edge_handles.artists)
 
+    def _init_shape(self, **props):
+        return Rectangle((0, 0), 0, 1, visible=False,
+                         rotate_around_center=True, **props)
+
     def _press(self, event):
         """Button press event handler."""
         # make the drawn box/line visible get the click-coordinates,
@@ -3017,9 +3025,17 @@ def _onmove(self, event):
             refx = event.xdata / (self._eventpress.xdata + 1e-6)
             refy = event.ydata / (self._eventpress.ydata + 1e-6)
 
+
+        x0, x1, y0, y1 = self._extents_on_press
         # resize an existing shape
-        if self._active_handle and self._active_handle != 'C':
-            x0, x1, y0, y1 = self._extents_on_press
+        if 'rotate' in state and self._active_handle in self._corner_order:
+            # calculate angle abc
+            a = np.array([self._eventpress.xdata, self._eventpress.ydata])
+            b = np.array(self.center)
+            c = np.array([event.xdata, event.ydata])
+            self._rotation = (np.arctan2(c[1]-b[1], c[0]-b[0]) -
+                              np.arctan2(a[1]-b[1], a[0]-b[0]))
+        elif self._active_handle and self._active_handle != 'C':
             sizepress = [x1 - x0, y1 - y0]
             center = [x0 + sizepress[0] / 2, y0 + sizepress[1] / 2]
 
@@ -3084,6 +3100,7 @@ def _onmove(self, event):
 
         # new shape
         else:
+            self._rotation = 0
             # Don't create a new rectangle if there is already one when
             # ignore_event_outside=True
             if self.ignore_event_outside and self._selection_completed:
@@ -3118,24 +3135,25 @@ def _onmove(self, event):
     @property
     def _rect_bbox(self):
         if self._drawtype == 'box':
-            x0 = self._selection_artist.get_x()
-            y0 = self._selection_artist.get_y()
-            width = self._selection_artist.get_width()
-            height = self._selection_artist.get_height()
-            return x0, y0, width, height
+            return self._selection_artist.get_bbox().bounds
         else:
             x, y = self._selection_artist.get_data()
             x0, x1 = min(x), max(x)
             y0, y1 = min(y), max(y)
             return x0, y0, x1 - x0, y1 - y0
 
+    def _get_rotation_transform(self):
+        return Affine2D().rotate_around(*self.center, self._rotation)
+
     @property
     def corners(self):
         """Corners of rectangle from lower left, moving clockwise."""
         x0, y0, width, height = self._rect_bbox
         xc = x0, x0 + width, x0 + width, x0
         yc = y0, y0, y0 + height, y0 + height
-        return xc, yc
+        transform = self._get_rotation_transform()
+        coords = transform.transform(np.array([xc, yc]).T).T
+        return coords[0], coords[1]
 
     @property
     def edge_centers(self):
@@ -3145,7 +3163,9 @@ def edge_centers(self):
         h = height / 2.
         xe = x0, x0 + w, x0 + width, x0 + w
         ye = y0 + h, y0, y0 + h, y0 + height
-        return xe, ye
+        transform = self._get_rotation_transform()
+        coords = transform.transform(np.array([xe, ye]).T).T
+        return coords[0], coords[1]
 
     @property
     def center(self):
@@ -3155,7 +3175,10 @@ def center(self):
 
     @property
     def extents(self):
-        """Return (xmin, xmax, ymin, ymax)."""
+        """
+        Return (xmin, xmax, ymin, ymax) as defined by the bounding box before
+        rotation.
+        """
         x0, y0, width, height = self._rect_bbox
         xmin, xmax = sorted([x0, x0 + width])
         ymin, ymax = sorted([y0, y0 + height])
@@ -3173,6 +3196,17 @@ def extents(self, extents):
         self.set_visible(self.visible)
         self.update()
 
+    @property
+    def rotation(self):
+        """Rotation in degree."""
+        return np.rad2deg(self._rotation)
+
+    @rotation.setter
+    def rotation(self, value):
+        self._rotation = np.deg2rad(value)
+        # call extents setter to draw shape and update handles positions
+        self.extents = self.extents
+
     draw_shape = _api.deprecate_privatize_attribute('3.5')
 
     def _draw_shape(self, extents):
@@ -3192,6 +3226,7 @@ def _draw_shape(self, extents):
             self._selection_artist.set_y(ymin)
             self._selection_artist.set_width(xmax - xmin)
             self._selection_artist.set_height(ymax - ymin)
+            self._selection_artist.set_angle(self.rotation)
 
         elif self._drawtype == 'line':
             self._selection_artist.set_data([xmin, xmax], [ymin, ymax])
@@ -3264,9 +3299,11 @@ class EllipseSelector(RectangleSelector):
     :doc:`/gallery/widgets/rectangle_selector`
     """
 
-    _shape_klass = Ellipse
     draw_shape = _api.deprecate_privatize_attribute('3.5')
 
+    def _init_shape(self, **props):
+        return Ellipse((0, 0), 0, 1, visible=False, **props)
+
     def _draw_shape(self, extents):
         x0, x1, y0, y1 = extents
         xmin, xmax = sorted([x0, x1])
@@ -3279,6 +3316,7 @@ def _draw_shape(self, extents):
             self._selection_artist.center = center
             self._selection_artist.width = 2 * a
             self._selection_artist.height = 2 * b
+            self._selection_artist.set_angle(self.rotation)
         else:
             rad = np.deg2rad(np.arange(31) * 12)
             x = a * np.cos(rad) + center[0]
@@ -3457,7 +3495,8 @@ def __init__(self, ax, onselect, useblit=False,
                                    move_all='shift', move='not-applicable',
                                    square='not-applicable',
                                    center='not-applicable',
-                                   data_coordinates='not-applicable')
+                                   data_coordinates='not-applicable',
+                                   rotate='not-applicable')
         super().__init__(ax, onselect, useblit=useblit,
                          state_modifier_keys=state_modifier_keys)
 
