Update lib/mpl_toolkits/mplot3d/axes3d.py

jgehrcke · dmcdougall · commit c50411e670ea · 2012-09-24T16:29:31.000+01:00
Axes3D.scatter: raise exception in case z.size neither being 1 nor equaling x.size and y.size (according to documentation).
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1937,8 +1937,11 @@ def scatter(self, xs, ys, zs=0, zdir='z', s=20, c='b', *args, **kwargs):
         zs = np.ma.ravel(zs)
         if xs.size != ys.size:
             raise ValueError("x and y must be the same size")
-        if xs.size != zs.size and zs.size == 1:
-            zs = np.array(zs[0] * xs.size)
+        if xs.size != zs.size:
+            if zs.size == 1:
+                zs = np.array(zs[0] * xs.size)
+            else:
+                raise ValueError("z must be either of same size as x and y or of size 1")
 
         s = np.ma.ravel(s)  # This doesn't have to match x, y in size.
 
