fixing quiver logic

neok-m4700 · neok-m4700 · commit 47825e924119 · 2017-05-29T02:04:14.000+02:00
added 2 simple tests
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -626,19 +626,19 @@ def _angles_lengths(self, U, V, eps=1):
 
     def _make_verts(self, U, V, angles):
         uv = (U + V * 1j)
-        str_angles = isinstance(angles, six.string_types)
-        if str_angles and (angles == 'xy' and self.scale_units == 'xy'):
+        str_angles = angles if isinstance(angles, six.string_types) else ''
+        if str_angles == 'xy' and self.scale_units == 'xy':
             # Here eps is 1 so that if we get U, V by diffing
             # the X, Y arrays, the vectors will connect the
             # points, regardless of the axis scaling (including log).
             angles, lengths = self._angles_lengths(U, V, eps=1)
-        elif str_angles and (angles == 'xy' or self.scale_units == 'xy'):
+        elif str_angles == 'xy' or self.scale_units == 'xy':
             # Calculate eps based on the extents of the plot
             # so that we don't end up with roundoff error from
             # adding a small number to a large.
             eps = np.abs(self.ax.dataLim.extents).max() * 0.001
             angles, lengths = self._angles_lengths(U, V, eps=eps)
-        if self.scale_units == 'xy':
+        if str_angles and self.scale_units == 'xy':
             a = lengths
         else:
             a = np.abs(uv)
@@ -665,9 +665,9 @@ def _make_verts(self, U, V, angles):
                 self.scale = scale * widthu_per_lenu
         length = a * (widthu_per_lenu / (self.scale * self.width))
         X, Y = self._h_arrows(length)
-        if str_angles and (angles == 'xy'):
+        if str_angles == 'xy':
             theta = angles
-        elif str_angles and (angles == 'uv'):
+        elif str_angles == 'uv':
             theta = np.angle(uv)
         else:
             theta = ma.masked_invalid(np.deg2rad(angles)).filled(0)
diff --git a/lib/matplotlib/tests/baseline_images/test_quiver/quiver_xy.png b/lib/matplotlib/tests/baseline_images/test_quiver/quiver_xy.png
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -155,6 +155,26 @@ def test_bad_masked_sizes():
         ax.barbs(x, y, u, v)
 
 
+def test_angles_and_scale():
+    # angles array + scale_units kwarg
+    fig, ax = plt.subplots()
+    X, Y = np.meshgrid(np.arange(15), np.arange(10))
+    U = V = np.ones_like(X)
+    phi = (np.random.rand(15, 10) - .5) * 150
+    ax.quiver(X, Y, U, V, angles=phi, scale_units='xy')
+
+
+@image_comparison(baseline_images=['quiver_xy'],
+                  extensions=['png'], remove_text=True)
+def test_quiver_xy():
+    # simple arrow pointing from SW to NE
+    fig, ax = plt.subplots(subplot_kw=dict(aspect='equal'))
+    ax.quiver(0, 0, 1, 1, angles='xy', scale_units='xy', scale=1)
+    ax.set_xlim(0, 1.1)
+    ax.set_ylim(0, 1.1)
+    ax.grid()
+
+
 def test_quiverkey_angles():
     # Check that only a single arrow is plotted for a quiverkey when an array
     # of angles is given to the original quiver plot
