Add set_offsets to quiver and make the attribute (N and XY) properties to avoid inconsistent state of quiver

ericpre · ericpre · commit 243af84b419f · 2023-09-05T21:36:57.000+01:00
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -449,7 +449,7 @@ class Quiver(mcollections.PolyCollection):
     The API methods are set_UVC(), set_U(), set_V() and set_C(), which
     can be used to change the size, orientation, and color of the
     arrows; their locations are fixed when the class is
-    instantiated.  Possibly this method will be useful
+    instantiated.  Possibly these methods will be useful
     in animations.
 
     Much of the work in this class is done in the draw()
@@ -477,8 +477,6 @@ def __init__(self, ax, *args,
         X, Y, U, V, C = _parse_args(*args, caller_name='quiver')
         self.X = X
         self.Y = Y
-        self.XY = np.column_stack((X, Y))
-        self.N = len(X)
         self.scale = scale
         self.headwidth = headwidth
         self.headlength = float(headlength)
@@ -525,6 +523,14 @@ def _init(self):
 
             self._dpi_at_last_init = self.axes.figure.dpi
 
+    @property
+    def N(self):
+        return len(self.X)
+
+    @property
+    def XY(self):
+        return np.column_stack((self.X, self.Y))
+
     def get_datalim(self, transData):
         trans = self.get_transform()
         offset_trf = self.get_offset_transform()
@@ -563,10 +569,13 @@ def set_UVC(self, U, V, C=None):
         ----------
         U : ArrayLike | None
             The x direction components of the arrows. If None it is unchanged.
+            The size must the same as the existing U, V or be one.
         V : ArrayLike | None
             The y direction components of the arrows. If None it is unchanged.
+            The size must the same as the existing U, V or be one.
         C : ArrayLike | None, optional
             The arrow colors. The default is None.
+            The size must the same as the existing U, V or be one.
         """
         if U is None:
             U = self.U
@@ -598,6 +607,19 @@ def set_UVC(self, U, V, C=None):
             self.set_array(C)
         self.stale = True
 
+    def set_offsets(self, xy):
+        """
+        Set the offsets for the arrows.  This saves the offsets passed
+        in and masks them as appropriate for the existing X/Y data.
+
+        Parameters
+        ----------
+        xy : sequence of pairs of floats
+        """
+        self.X, self.Y = xy[:, 0], xy[:, 1]
+        super().set_offsets(xy)
+        self.stale = True
+
     def _dots_per_unit(self, units):
         """Return a scale factor for converting from units to pixels."""
         bb = self.axes.bbox
diff --git a/lib/matplotlib/quiver.pyi b/lib/matplotlib/quiver.pyi
@@ -54,11 +54,9 @@ class QuiverKey(martist.Artist):
 class Quiver(mcollections.PolyCollection):
     X: ArrayLike
     Y: ArrayLike
-    XY: ArrayLike
     U: ArrayLike
     V: ArrayLike
     Umask: ArrayLike
-    N: int
     scale: float | None
     headwidth: float
     headlength: float
@@ -121,13 +119,18 @@ class Quiver(mcollections.PolyCollection):
         pivot: Literal["tail", "mid", "middle", "tip"] = ...,
         **kwargs
     ) -> None: ...
+    @property
+    def N(self) -> int: ...
+    @property
+    def XY(self) -> ArrayLike: ...
     def get_datalim(self, transData: Transform) -> Bbox: ...
     def set_U(self, U: ArrayLike) -> None: ...
     def set_V(self, V: ArrayLike) -> None: ...
     def set_C(self, C: ArrayLike) -> None: ...
     def set_UVC(
         self, U: ArrayLike | None, V: ArrayLike | None, C: ArrayLike | None = ...
     ) -> None: ...
+    def set_offests(self, xy: ArrayLike) -> None: ...
     @property
     def quiver_doc(self) -> str: ...
 
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -366,6 +366,9 @@ def test_quiver_offsets():
     qc.set_offsets(new_offsets)
 
     np.testing.assert_allclose(qc.get_offsets(), new_offsets)
+    np.testing.assert_allclose(qc.X, new_offsets[::, 0])
+    np.testing.assert_allclose(qc.Y, new_offsets[::, 1])
+    np.testing.assert_allclose(qc.XY, new_offsets)
 
 
 def test_quiver_UVC():
