Fix positional/kwarg handling of the Z argument

WeatherGod · WeatherGod · commit 7ef9a89e8967 · 2014-02-28T08:49:27.000-05:00
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -1827,12 +1827,17 @@ def plot_trisurf(self, *args, **kwargs):
         lightsource = kwargs.pop('lightsource', None)
 
         tri, args, kwargs = Triangulation.get_from_args_and_kwargs(*args, **kwargs)
-        z = np.asarray(args[0])
+        if 'Z' in kwargs:
+            z = np.asarray(kwargs.pop('Z'))
+        else:
+            z = np.asarray(args[0])
+            # We do this so Z doesn't get passed as an arg to PolyCollection
+            args = args[1:]
 
         triangles = tri.get_masked_triangles()
-        xt = tri.x[triangles][...,np.newaxis]
-        yt = tri.y[triangles][...,np.newaxis]
-        zt = np.array(z)[triangles][...,np.newaxis]
+        xt = tri.x[triangles][..., np.newaxis]
+        yt = tri.y[triangles][..., np.newaxis]
+        zt = z[triangles][..., np.newaxis]
 
         verts = np.concatenate((xt, yt, zt), axis=2)
 
@@ -2024,10 +2029,12 @@ def tricontour(self, *args, **kwargs):
                 *args, **kwargs)
         X = tri.x
         Y = tri.y
-        Z = args[0]
-
-        # We do this so Z doesn't get passed as an arg to Axes.tricontour
-        args = args[1:]
+        if 'Z' in kwargs:
+            Z = kwargs.pop('Z')
+        else:
+            Z = args[0]
+            # We do this so Z doesn't get passed as an arg to Axes.tricontour
+            args = args[1:]
 
         jX, jY, jZ = art3d.rotate_axes(X, Y, Z, zdir)
         tri = Triangulation(jX, jY, tri.triangles, tri.mask)
@@ -2109,10 +2116,12 @@ def tricontourf(self, *args, **kwargs):
                 *args, **kwargs)
         X = tri.x
         Y = tri.y
-        Z = args[0]
-
-        # We do this so Z doesn't get passed as an arg to Axes.tricontour
-        args = args[1:]
+        if 'Z' in kwargs:
+            Z = kwargs.pop('Z')
+        else:
+            Z = args[0]
+            # We do this so Z doesn't get passed as an arg to Axes.tricontourf
+            args = args[1:]
 
         jX, jY, jZ = art3d.rotate_axes(X, Y, Z, zdir)
         tri = Triangulation(jX, jY, tri.triangles, tri.mask)
