Improved checking logic of _check_xyz in contour.py, added tests for exceptions

astrofrog · astrofrog · commit 4910610d86e1 · 2012-11-13T12:18:52.000+01:00
diff --git a/lib/matplotlib/__init__.py b/lib/matplotlib/__init__.py
@@ -1077,6 +1077,7 @@ def tk_window_focus():
     'matplotlib.tests.test_cbook',
     'matplotlib.tests.test_colorbar',
     'matplotlib.tests.test_colors',
+    'matplotlib.tests.test_contour',
     'matplotlib.tests.test_dates',
     'matplotlib.tests.test_delaunay',
     'matplotlib.tests.test_figure',
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -1424,20 +1424,40 @@ def _check_xyz(self, args, kwargs):
         x = np.asarray(x, dtype=np.float64)
         y = np.asarray(y, dtype=np.float64)
         z = ma.asarray(args[2], dtype=np.float64)
+
         if z.ndim != 2:
             raise TypeError("Input z must be a 2D array.")
         else:
             Ny, Nx = z.shape
-        if x.shape == z.shape and y.shape == z.shape:
-            return x, y, z
-        if x.ndim != 1 or y.ndim != 1:
+
+        if x.ndim != y.ndim:
+            raise TypeError("Number of dimensions of x and y should match.")
+
+        if x.ndim == 1:
+
+            nx, = x.shape
+            ny, = y.shape
+
+            if nx != Nx:
+                raise TypeError("Length of x must be number of columns in z.")
+
+            if ny != Ny:
+                raise TypeError("Length of y must be number of rows in z.")
+
+            x, y = np.meshgrid(x, y)
+
+        elif x.ndim == 2:
+
+            if x.shape != z.shape:
+                raise TypeError("Shape of x does not match that of z: found {0} instead of {1}.".format(x.shape, z.shape))
+
+            if y.shape != z.shape:
+                raise TypeError("Shape of y does not match that of z: found {0} instead of {1}.".format(y.shape, z.shape))
+
+        else:
+
             raise TypeError("Inputs x and y must be 1D or 2D.")
-        nx, = x.shape
-        ny, = y.shape
-        if nx != Nx or ny != Ny:
-            raise TypeError("Length of x must be number of columns in z,\n" +
-                            "and length of y must be number of rows.")
-        x, y = np.meshgrid(x, y)
+
         return x, y, z
 
     def _initialize_x_y(self, z):
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -0,0 +1,127 @@
+import numpy as np
+from matplotlib import pyplot as plt
+
+
+def test_contour_shape_1d_valid():
+
+    x = np.arange(10)
+    y = np.arange(9)
+    z = np.random.random((9, 10))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.contour(x, y, z)
+
+
+def test_contour_shape_2d_valid():
+
+    x = np.arange(10)
+    y = np.arange(9)
+    xg, yg = np.meshgrid(x, y)
+    z = np.random.random((9, 10))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+    ax.contour(xg, yg, z)
+
+
+def test_contour_shape_mismatch_1():
+
+    x = np.arange(9)
+    y = np.arange(9)
+    z = np.random.random((9, 10))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    try:
+        ax.contour(x, y, z)
+    except TypeError as exc:
+        assert exc.args[0] == 'Length of x must be number of columns in z.'
+
+
+def test_contour_shape_mismatch_2():
+
+    x = np.arange(10)
+    y = np.arange(10)
+    z = np.random.random((9, 10))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    try:
+        ax.contour(x, y, z)
+    except TypeError as exc:
+        assert exc.args[0] == 'Length of y must be number of rows in z.'
+
+
+def test_contour_shape_mismatch_3():
+
+    x = np.arange(10)
+    y = np.arange(10)
+    xg, yg = np.meshgrid(x, y)
+    z = np.random.random((9, 10))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    try:
+        ax.contour(xg, y, z)
+    except TypeError as exc:
+        assert exc.args[0] == 'Number of dimensions of x and y should match.'
+
+    try:
+        ax.contour(x, yg, z)
+    except TypeError as exc:
+        assert exc.args[0] == 'Number of dimensions of x and y should match.'
+
+
+def test_contour_shape_mismatch_4():
+
+    g = np.random.random((9, 10))
+    b = np.random.random((9, 9))
+    z = np.random.random((9, 10))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    try:
+        ax.contour(b, g, z)
+    except TypeError as exc:
+        print exc.args[0]
+        assert exc.args[0] == 'Shape of x does not match that of z: found (9, 9) instead of (9, 10).'
+
+    try:
+        ax.contour(g, b, z)
+    except TypeError as exc:
+        assert exc.args[0] == 'Shape of y does not match that of z: found (9, 9) instead of (9, 10).'
+
+
+def test_contour_shape_invalid_1():
+
+    x = np.random.random((3, 3, 3))
+    y = np.random.random((3, 3, 3))
+    z = np.random.random((9, 10))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    try:
+        ax.contour(x, y, z)
+    except TypeError as exc:
+        assert exc.args[0] == 'Inputs x and y must be 1D or 2D.'
+
+
+def test_contour_shape_invalid_2():
+
+    x = np.random.random((3, 3, 3))
+    y = np.random.random((3, 3, 3))
+    z = np.random.random((3, 3, 3))
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    try:
+        ax.contour(x, y, z)
+    except TypeError as exc:
+        assert exc.args[0] == 'Input z must be a 2D array.'
