Take into account aspect ratio in square state

ericpre · ericpre · commit 034c251d5fa8 · 2021-08-19T15:51:06.000+01:00
diff --git a/lib/matplotlib/tests/test_widgets.py b/lib/matplotlib/tests/test_widgets.py
@@ -161,6 +161,7 @@ def onselect(epress, erelease):
     tool.add_default_state('move')
     tool.add_default_state('square')
     tool.add_default_state('center')
+    tool.add_default_state('data_coordinates')
 
 
 @pytest.mark.parametrize('use_default_state', [True, False])
@@ -377,6 +378,42 @@ def onselect(epress, erelease):
                             ydata_new, extents[3] - ydiff)
 
 
+def test_rectangle_resize_square_center_aspect():
+    ax = get_ax()
+    ax.set_aspect(0.8)
+
+    def onselect(epress, erelease):
+        pass
+
+    tool = widgets.RectangleSelector(ax, onselect, interactive=True)
+    # Create rectangle
+    _resize_rectangle(tool, 70, 65, 120, 115)
+    tool.add_default_state('square')
+    tool.add_default_state('center')
+    assert tool.extents == (70.0, 120.0, 65.0, 115.0)
+
+    # resize E handle
+    extents = tool.extents
+    xdata, ydata = extents[1], extents[3]
+    xdiff = 10
+    xdata_new, ydata_new = xdata + xdiff, ydata
+    _resize_rectangle(tool, xdata, ydata, xdata_new, ydata_new)
+    assert_allclose(tool.extents, [extents[0] - xdiff, xdata_new,
+                                   46.25, 133.75])
+
+    # use data coordinates
+    do_event(tool, 'on_key_press', key='d')
+    # resize E handle
+    extents = tool.extents
+    xdata, ydata, width = extents[1], extents[3], extents[1] - extents[0]
+    xdiff, ycenter = 10,  extents[2] + (extents[3] - extents[2]) / 2
+    xdata_new, ydata_new = xdata + xdiff, ydata
+    ychange = width / 2 + xdiff
+    _resize_rectangle(tool, xdata, ydata, xdata_new, ydata_new)
+    assert_allclose(tool.extents, [extents[0] - xdiff, xdata_new,
+                                   ycenter - ychange, ycenter + ychange])
+
+
 def test_ellipse():
     """For ellipse, test out the key modifiers"""
     ax = get_ax()
@@ -415,7 +452,7 @@ def onselect(epress, erelease):
     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]
+    assert extents == [10, 35, 10, 35]
 
     # create a square from center
     do_event(tool, 'on_key_press', xdata=100, ydata=100, button=1,
@@ -426,7 +463,7 @@ def onselect(epress, erelease):
     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]
+    assert extents == [70, 130, 70, 130]
 
     assert tool.geometry.shape == (2, 73)
     assert_allclose(tool.geometry[:, 0], [70., 100])
diff --git a/lib/matplotlib/widgets.py b/lib/matplotlib/widgets.py
@@ -1796,7 +1796,8 @@ def __init__(self, ax, onselect, useblit=False, button=None,
         self.connect_default_events()
 
         self._state_modifier_keys = dict(move=' ', clear='escape',
-                                         square='shift', center='control')
+                                         square='shift', center='control',
+                                         data_coordinates='d')
         self._state_modifier_keys.update(state_modifier_keys or {})
 
         self.background = None
@@ -1985,6 +1986,12 @@ 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, modifier) in self._state_modifier_keys.items():
                 if modifier in key:
                     self._state.add(state)
@@ -2152,7 +2159,8 @@ def __init__(self, ax, onselect, direction, minspan=0, useblit=False,
         if state_modifier_keys is None:
             state_modifier_keys = dict(clear='escape',
                                        square='not-applicable',
-                                       center='not-applicable')
+                                       center='not-applicable',
+                                       data_coordinates='not-applicable')
         super().__init__(ax, onselect, useblit=useblit, button=button,
                          state_modifier_keys=state_modifier_keys)
 
@@ -2713,8 +2721,12 @@ def onselect(eclick: MouseEvent, erelease: MouseEvent)
         - "clear": Clear the current shape, default: "escape".
         - "square": Makes the shape square, default: "shift".
         - "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"
 
-        "square" and "center" can be combined.
+        "square" and "center" can be combined. The square shape can be defined
+        in data or figure coordinates as determined by the ``data_coordinates``
+        modifier, which can be enable and disable by pressing the 'd' key.
 
     drag_from_anywhere : bool, default: False
         If `True`, the widget can be moved by clicking anywhere within
@@ -2947,61 +2959,75 @@ def _onmove(self, event):
         """Motion notify event handler."""
 
         state = self._state | self._default_state
+
+        dx = event.xdata - self._eventpress.xdata
+        dy = event.ydata - self._eventpress.ydata
+        refmax = None
+        if 'data_coordinates' in state:
+            aspect_ratio = 1
+            refx, refy = dx, dy
+        else:
+            figure_size = self.ax.get_figure().get_size_inches()
+            ll, ur = self.ax.get_position() * figure_size
+            width, height = ur - ll
+            aspect_ratio = height / width * self.ax.get_data_ratio()
+            refx = event.xdata / (self._eventpress.xdata + 1e-6)
+            refy = event.ydata / (self._eventpress.ydata + 1e-6)
+
         # resize an existing shape
         if self._active_handle and self._active_handle != 'C':
             x0, x1, y0, y1 = self._extents_on_press
             sizepress = [x1 - x0, y1 - y0]
             center = [x0 + sizepress[0] / 2, y0 + sizepress[1] / 2]
-            dx = event.xdata - self._eventpress.xdata
-            dy = event.ydata - self._eventpress.ydata
-
-            # change sign of relative changes to simplify calculation
-            # Switch variables so that only x1 and/or y1 are updated on move
-            x_factor = y_factor = 1
-            if 'W' in self._active_handle:
-                x_factor *= -1
-                dx *= x_factor
-                x0 = x1
-            if 'S' in self._active_handle:
-                y_factor *= -1
-                dy *= y_factor
-                y0 = y1
 
             # from center
             if 'center' in state:
                 if 'square' in state:
-                    if self._active_handle in ['E', 'W']:
-                        # using E, W handle we need to update dy accordingly
-                        dy = dx
-                    elif self._active_handle in ['S', 'N']:
-                        # using S, N handle, we need to update dx accordingly
-                        dx = dy
+                    # when using a corner, find which reference to use
+                    if self._active_handle in self._corner_order:
+                        refmax = max(refx, refy, key=abs)
+                    if self._active_handle in ['E', 'W'] or refmax == refx:
+                        dw = event.xdata - center[0]
+                        dh = dw / aspect_ratio
                     else:
-                        dx = dy = max(dx, dy, key=abs)
-
-                dw = sizepress[0] / 2 + dx
-                dh = sizepress[1] / 2 + dy
-
-                if 'square' not in state:
+                        dh = event.ydata - center[1]
+                        dw = dh * aspect_ratio
+                else:
+                    dw = sizepress[0] / 2
+                    dh = sizepress[1] / 2
                     # cancel changes in perpendicular direction
-                    if self._active_handle in ['E', 'W']:
-                        dh = sizepress[1] / 2
-                    if self._active_handle in ['N', 'S']:
-                        dw = sizepress[0] / 2
+                    if self._active_handle in ['E', 'W'] + self._corner_order:
+                        dw = abs(event.xdata - center[0])
+                    if self._active_handle in ['N', 'S'] + self._corner_order:
+                        dh = abs(event.ydata - center[1])
 
                 x0, x1, y0, y1 = (center[0] - dw, center[0] + dw,
                                   center[1] - dh, center[1] + dh)
 
             else:
+                # change sign of relative changes to simplify calculation
+                # Switch variables so that x1 and/or y1 are updated on move
+                x_factor = y_factor = 1
+                if 'W' in self._active_handle:
+                    x0 = x1
+                    x_factor *= -1
+                if 'S' in self._active_handle:
+                    y0 = y1
+                    y_factor *= -1
+                if self._active_handle in ['E', 'W'] + self._corner_order:
+                    x1 = event.xdata
+                if self._active_handle in ['N', 'S'] + self._corner_order:
+                    y1 = event.ydata
                 if 'square' in state:
-                    dx = dy = max(dx, dy, key=abs)
-                    x1 = x0 + x_factor * (dx + sizepress[0])
-                    y1 = y0 + y_factor * (dy + sizepress[1])
-                else:
-                    if self._active_handle in ['E', 'W'] + self._corner_order:
-                        x1 = event.xdata
-                    if self._active_handle in ['N', 'S'] + self._corner_order:
-                        y1 = event.ydata
+                    # when using a corner, find which reference to use
+                    if self._active_handle in self._corner_order:
+                        refmax = max(refx, refy, key=abs)
+                    if self._active_handle in ['E', 'W'] or refmax == refx:
+                        sign = np.sign(event.ydata - y0)
+                        y1 = y0 + sign * abs(x1 - x0) / aspect_ratio
+                    else:
+                        sign = np.sign(event.xdata - x0)
+                        x1 = x0 + sign * abs(y1 - y0) * aspect_ratio
 
         # move existing shape
         elif self._active_handle == 'C':
@@ -3020,21 +3046,16 @@ def _onmove(self, event):
             if self.ignore_event_outside and self._selection_completed:
                 return
             center = [self._eventpress.xdata, self._eventpress.ydata]
-            center_pix = [self._eventpress.x, self._eventpress.y]
             dx = (event.xdata - center[0]) / 2.
             dy = (event.ydata - center[1]) / 2.
 
             # square shape
             if 'square' in state:
-                dx_pix = abs(event.x - center_pix[0])
-                dy_pix = abs(event.y - center_pix[1])
-                if not dx_pix:
-                    return
-                maxd = max(abs(dx_pix), abs(dy_pix))
-                if abs(dx_pix) < maxd:
-                    dx *= maxd / (abs(dx_pix) + 1e-6)
-                if abs(dy_pix) < maxd:
-                    dy *= maxd / (abs(dy_pix) + 1e-6)
+                refmax = max(refx, refy, key=abs)
+                if refmax == refx:
+                    dy = dx / aspect_ratio
+                else:
+                    dx = dy * aspect_ratio
 
             # from center
             if 'center' in state:
@@ -3390,7 +3411,8 @@ def __init__(self, ax, onselect, useblit=False,
         state_modifier_keys = dict(clear='escape', move_vertex='control',
                                    move_all='shift', move='not-applicable',
                                    square='not-applicable',
-                                   center='not-applicable')
+                                   center='not-applicable',
+                                   data_coordinates='not-applicable')
         super().__init__(ax, onselect, useblit=useblit,
                          state_modifier_keys=state_modifier_keys)
 
