Merge pull request #18807 from ianhi/FancyArrow-set-positions

timhoffm · web-flow · commit 5120b1956bdb · 2021-04-24T01:31:32.000+02:00
make FancyArrow animatable
diff --git a/doc/users/next_whats_new/animatable_FancyArrow.rst b/doc/users/next_whats_new/animatable_FancyArrow.rst
@@ -0,0 +1,5 @@
+``set_data`` method for ``FancyArrow`` patch
+--------------------------------------------
+
+`.FancyArrow`, the patch returned by ``ax.arrow``, now has a ``set_data``
+method that allows for animating the arrow.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5032,12 +5032,6 @@ def arrow(self, x, y, dx, dy, **kwargs):
 
         Parameters
         ----------
-        x, y : float
-            The x and y coordinates of the arrow base.
-
-        dx, dy : float
-            The length of the arrow along x and y direction.
-
         %(FancyArrow)s
 
         Returns
diff --git a/lib/matplotlib/patches.py b/lib/matplotlib/patches.py
@@ -1322,6 +1322,12 @@ def __init__(self, x, y, dx, dy, width=0.001, length_includes_head=False,
         """
         Parameters
         ----------
+        x, y : float
+            The x and y coordinates of the arrow base.
+
+        dx, dy : float
+            The length of the arrow along x and y direction.
+
         width : float, default: 0.001
             Width of full arrow tail.
 
@@ -1350,22 +1356,82 @@ def __init__(self, x, y, dx, dy, width=0.001, length_includes_head=False,
 
             %(Patch_kwdoc)s
         """
-        if head_width is None:
-            head_width = 3 * width
-        if head_length is None:
+        self._x = x
+        self._y = y
+        self._dx = dx
+        self._dy = dy
+        self._width = width
+        self._length_includes_head = length_includes_head
+        self._head_width = head_width
+        self._head_length = head_length
+        self._shape = shape
+        self._overhang = overhang
+        self._head_starts_at_zero = head_starts_at_zero
+        self._make_verts()
+        super().__init__(self.verts, closed=True, **kwargs)
+
+    def set_data(self, *, x=None, y=None, dx=None, dy=None, width=None,
+                 head_width=None, head_length=None):
+        """
+        Set `.FancyArrow` x, y, dx, dy, width, head_with, and head_length.
+        Values left as None will not be updated.
+
+        Parameters
+        ----------
+        x, y : float or None, default: None
+            The x and y coordinates of the arrow base.
+
+        dx, dy : float or None, default: None
+            The length of the arrow along x and y direction.
+
+        width: float or None, default: None
+            Width of full arrow tail.
+
+        head_width: float or None, default: None
+            Total width of the full arrow head.
+
+        head_length: float or None, default: None
+            Length of arrow head.
+        """
+        if x is not None:
+            self._x = x
+        if y is not None:
+            self._y = y
+        if dx is not None:
+            self._dx = dx
+        if dy is not None:
+            self._dy = dy
+        if width is not None:
+            self._width = width
+        if head_width is not None:
+            self._head_width = head_width
+        if head_length is not None:
+            self._head_length = head_length
+        self._make_verts()
+        self.set_xy(self.verts)
+
+    def _make_verts(self):
+        if self._head_width is None:
+            head_width = 3 * self._width
+        else:
+            head_width = self._head_width
+        if self._head_length is None:
             head_length = 1.5 * head_width
+        else:
+            head_length = self._head_length
 
-        distance = np.hypot(dx, dy)
+        distance = np.hypot(self._dx, self._dy)
 
-        if length_includes_head:
+        if self._length_includes_head:
             length = distance
         else:
             length = distance + head_length
         if not length:
-            verts = np.empty([0, 2])  # display nothing if empty
+            self.verts = np.empty([0, 2])  # display nothing if empty
         else:
             # start by drawing horizontal arrow, point at (0, 0)
-            hw, hl, hs, lw = head_width, head_length, overhang, width
+            hw, hl = head_width, head_length
+            hs, lw = self._overhang, self._width
             left_half_arrow = np.array([
                 [0.0, 0.0],                 # tip
                 [-hl, -hw / 2],             # leftmost
@@ -1374,36 +1440,37 @@ def __init__(self, x, y, dx, dy, width=0.001, length_includes_head=False,
                 [-length, 0],
             ])
             # if we're not including the head, shift up by head length
-            if not length_includes_head:
+            if not self._length_includes_head:
                 left_half_arrow += [head_length, 0]
             # if the head starts at 0, shift up by another head length
-            if head_starts_at_zero:
+            if self._head_starts_at_zero:
                 left_half_arrow += [head_length / 2, 0]
             # figure out the shape, and complete accordingly
-            if shape == 'left':
+            if self._shape == 'left':
                 coords = left_half_arrow
             else:
                 right_half_arrow = left_half_arrow * [1, -1]
-                if shape == 'right':
+                if self._shape == 'right':
                     coords = right_half_arrow
-                elif shape == 'full':
+                elif self._shape == 'full':
                     # The half-arrows contain the midpoint of the stem,
                     # which we can omit from the full arrow. Including it
                     # twice caused a problem with xpdf.
                     coords = np.concatenate([left_half_arrow[:-1],
                                              right_half_arrow[-2::-1]])
                 else:
-                    raise ValueError("Got unknown shape: %s" % shape)
+                    raise ValueError("Got unknown shape: %s" % self.shape)
             if distance != 0:
-                cx = dx / distance
-                sx = dy / distance
+                cx = self._dx / distance
+                sx = self._dy / distance
             else:
                 # Account for division by zero
                 cx, sx = 0, 1
             M = [[cx, sx], [-sx, cx]]
-            verts = np.dot(coords, M) + (x + dx, y + dy)
-
-        super().__init__(verts, closed=True, **kwargs)
+            self.verts = np.dot(coords, M) + [
+                self._x + self._dx,
+                self._y + self._dy,
+            ]
 
 
 docstring.interpd.update(
diff --git a/lib/matplotlib/tests/test_patches.py b/lib/matplotlib/tests/test_patches.py
@@ -548,7 +548,37 @@ def test_fancyarrow_units():
     dtime = datetime(2000, 1, 1)
     fig, ax = plt.subplots()
     arrow = FancyArrowPatch((0, dtime), (0.01, dtime))
-    ax.add_patch(arrow)
+
+
+def test_fancyarrow_setdata():
+    fig, ax = plt.subplots()
+    arrow = ax.arrow(0, 0, 10, 10, head_length=5, head_width=1, width=.5)
+    expected1 = np.array(
+      [[13.54, 13.54],
+       [10.35,  9.65],
+       [10.18,  9.82],
+       [0.18, -0.18],
+       [-0.18,  0.18],
+       [9.82, 10.18],
+       [9.65, 10.35],
+       [13.54, 13.54]]
+    )
+    assert np.allclose(expected1, np.round(arrow.verts, 2))
+
+    expected2 = np.array(
+      [[16.71, 16.71],
+       [16.71, 15.29],
+       [16.71, 15.29],
+       [1.71,  0.29],
+       [0.29,  1.71],
+       [15.29, 16.71],
+       [15.29, 16.71],
+       [16.71, 16.71]]
+    )
+    arrow.set_data(
+        x=1, y=1, dx=15, dy=15, width=2, head_width=2, head_length=1
+    )
+    assert np.allclose(expected2, np.round(arrow.verts, 2))
 
 
 @image_comparison(["large_arc.svg"], style="mpl20")
