ENH: Improve barbs() error message for badly-sized arrays.

dopplershift · dopplershift · commit b915ae17b447 · 2016-11-04T14:47:25.000-06:00
With this change, when given masked arrays with incompatible sizes,
barbs() will raise a ValueError with a message stating as much.
Previously, it would error with a TypeError that did not make it clear
at all the source of the problem.
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,
@@ -1161,6 +1169,7 @@ def set_offsets(self, xy):
         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()
