Merge pull request #7407 from dopplershift/barbs-error

NelleV · web-flow · commit dc1368842e05 · 2016-11-04T16:34:44.000-07:00
Improve barbs() error message
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -392,6 +392,12 @@ def _parse_args(*args):
     return X, Y, U, V, C
 
 
+def _check_consistent_shapes(*arrays):
+    all_shapes = set(a.shape for a in arrays)
+    if len(all_shapes) != 1:
+        raise ValueError('The shapes of the passed in arrays do not match.')
+
+
 class Quiver(mcollections.PolyCollection):
     """
     Specialized PolyCollection for arrows.
@@ -1124,9 +1130,11 @@ def set_UVC(self, U, V, C=None):
             x, y, u, v, c = delete_masked_points(self.x.ravel(),
                                                  self.y.ravel(),
                                                  self.u, self.v, c)
+            _check_consistent_shapes(x, y, u, v, c)
         else:
             x, y, u, v = delete_masked_points(self.x.ravel(), self.y.ravel(),
                                               self.u, self.v)
+            _check_consistent_shapes(x, y, u, v)
 
         magnitude = np.hypot(u, v)
         flags, barbs, halves, empty = self._find_tails(magnitude,
@@ -1151,16 +1159,17 @@ def set_UVC(self, U, V, C=None):
 
     def set_offsets(self, xy):
         """
-        Set the offsets for the barb polygons.  This saves the offets passed in
-        and actually sets version masked as appropriate for the existing U/V
-        data. *offsets* should be a sequence.
+        Set the offsets for the barb polygons.  This saves the offsets passed
+        in and actually sets version masked as appropriate for the existing
+        U/V data. *offsets* should be a sequence.
 
         ACCEPTS: sequence of pairs of floats
         """
         self.x = xy[:, 0]
         self.y = xy[:, 1]
         x, y, u, v = delete_masked_points(self.x.ravel(), self.y.ravel(),
                                           self.u, self.v)
+        _check_consistent_shapes(x, y, u, v)
         xy = np.hstack((x[:, np.newaxis], y[:, np.newaxis]))
         mcollections.PolyCollection.set_offsets(self, xy)
         self.stale = True
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -1,6 +1,7 @@
 from __future__ import print_function
 import warnings
 import numpy as np
+from nose.tools import raises
 import sys
 from matplotlib import pyplot as plt
 from matplotlib.testing.decorators import cleanup
@@ -133,6 +134,20 @@ def test_barbs():
              cmap='viridis')
 
 
+@cleanup
+@raises(ValueError)
+def test_bad_masked_sizes():
+    'Test error handling when given differing sized masked arrays'
+    x = np.arange(3)
+    y = np.arange(3)
+    u = np.ma.array(15. * np.ones((4,)))
+    v = np.ma.array(15. * np.ones_like(u))
+    u[1] = np.ma.masked
+    v[1] = np.ma.masked
+    fig, ax = plt.subplots()
+    ax.barbs(x, y, u, v)
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()
