Allow numbers to set uvc for all arrows in quiver.set_UVC, fixes #16743

maxnoe · maxnoe · commit a547f0d86dfb · 2020-03-12T15:31:58.000+01:00
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -585,7 +585,7 @@ def set_UVC(self, U, V, C=None):
         if C is not None:
             C = ma.masked_invalid(C, copy=True).ravel()
         for name, var in zip(('U', 'V', 'C'), (U, V, C)):
-            if var is not None and var.size != self.N:
+            if var is not None and var.size != self.N and var.size != 1:
                 raise ValueError(f'Argument {name} has a size {var.size}'
                                  f' which does not match {self.N},'
                                  ' the number of arrow positions')
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -247,3 +247,15 @@ def test_quiverkey_angles():
     # The arrows are only created when the key is drawn
     fig.canvas.draw()
     assert len(qk.verts) == 1
+
+
+def test_quiver_setuvc_numbers():
+    """Check that it is possible to set all arrow UVC to the same numbers"""
+
+    fig, ax = plt.subplots()
+
+    X, Y = np.meshgrid(np.arange(2), np.arange(2))
+    U = V = np.ones_like(X)
+
+    q = ax.quiver(X, Y, U, V)
+    q.set_UVC(0, 1)
