Commits

Brian Kearns committed 78f7681

add better tests for vstack/hstack/dstack and pep8-ify

  • Participants
  • Parent commits 5fb4c22

Comments (0)

Files changed (1)

File pypy/module/test_lib_pypy/numpypy/core/test_shape_base.py

 from pypy.module.micronumpy.test.test_base import BaseNumpyAppTest
 
+
 class AppTestShapeBase(BaseNumpyAppTest):
 
     def test_atleast_1d(self):
         from numpypy import array, array_equal
         import numpypy as np
         a = np.atleast_1d(1.0)
-        assert np.array_equal(a, [ 1.])
-        
-        x = np.arange(9.0).reshape(3,3)
+        assert np.array_equal(a, [1.])
+
+        x = np.arange(9.0).reshape(3, 3)
         a = np.atleast_1d(x)
-        assert np.array_equal(a, [[ 0.,  1.,  2.],
-                                  [ 3.,  4.,  5.],
-                                  [ 6.,  7.,  8.]])
+        assert np.array_equal(a, [[0.,  1.,  2.],
+                                  [3.,  4.,  5.],
+                                  [6.,  7.,  8.]])
         assert np.atleast_1d(x) is x
 
         a = np.atleast_1d(1, [3, 4])
     def test_atleast_2d(self):
         import numpypy as np
         a = np.atleast_2d(3.0)
-        assert np.array_equal(a, [[ 3.]])
+        assert np.array_equal(a, [[3.]])
 
         x = np.arange(3.0)
         a = np.atleast_2d(x)
-        assert np.array_equal(a, [[ 0.,  1.,  2.]])
+        assert np.array_equal(a, [[0., 1., 2.]])
 
         a = np.atleast_2d(1, [1, 2], [[1, 2]])
         assert len(a) == 3
         import numpypy as np
 
         a = np.atleast_3d(3.0)
-        assert np.array_equal(a, [[[ 3.]]])
+        assert np.array_equal(a, [[[3.]]])
 
         x = np.arange(3.0)
         assert np.atleast_3d(x).shape == (1, 3, 1)
 
-        x = np.arange(12.0).reshape(4,3)
+        x = np.arange(12.0).reshape(4, 3)
         assert np.atleast_3d(x).shape == (4, 3, 1)
 
         a = np.atleast_3d([1, 2])
 
         a = np.array([1, 2, 3])
         b = np.array([2, 3, 4])
-        c = np.vstack((a,b))
+        c = np.vstack((a, b))
         assert np.array_equal(c, [[1, 2, 3],
                                   [2, 3, 4]])
 
         a = np.array([[1], [2], [3]])
         b = np.array([[2], [3], [4]])
-        c = np.vstack((a,b))
+        c = np.vstack((a, b))
         assert np.array_equal(c, [[1],
                                   [2],
                                   [3],
                                   [3],
                                   [4]])
 
+        for shape1, shape2 in [[(2, 1), (3, 1)],
+                               [(2, 4), [3, 4]]]:
+            a, b = np.ones(shape1), np.ones(shape2)
+            assert np.all(np.vstack((a, b)) ==
+                          np.ones((a.shape[0] + b.shape[0],
+                                   a.shape[1])))
+
+        for shape1, shape2 in [[(3, 2, 4), (7, 2, 4)],
+                               [(0, 2, 7), (10, 2, 7)],
+                               [(0, 2, 7), (0, 2, 7)]]:
+            a, b = np.ones(shape1), np.ones(shape2)
+            assert np.all(np.vstack((a, b)) ==
+                          np.ones((a.shape[0] + b.shape[0],
+                                   a.shape[1],
+                                   a.shape[2])))
+
     def test_hstack(self):
         import numpypy as np
-        a = np.array((1,2,3))
-        b = np.array((2,3,4))
-        c = np.hstack((a,b))
+        a = np.array((1, 2, 3))
+        b = np.array((2, 3, 4))
+        c = np.hstack((a, b))
         assert np.array_equal(c, [1, 2, 3, 2, 3, 4])
-        
-        a = np.array([[1],[2],[3]])
-        b = np.array([[2],[3],[4]])
-        c = np.hstack((a,b))
+
+        a = np.array([[1], [2], [3]])
+        b = np.array([[2], [3], [4]])
+        c = np.hstack((a, b))
         assert np.array_equal(c, [[1, 2],
                                   [2, 3],
                                   [3, 4]])
 
+        for shape1, shape2 in [[(1, 2), (1, 3)],
+                               [(4, 2), (4, 3)]]:
+            a, b = np.ones(shape1), np.ones(shape2)
+            assert np.all(np.hstack((a, b)) ==
+                          np.ones((a.shape[0],
+                                   a.shape[1] + b.shape[1])))
+
+        for shape1, shape2 in [[(2, 3, 4), (2, 7, 4)],
+                               [(1, 4, 7), (1, 10, 7)],
+                               [(1, 4, 7), (1, 0, 7)],
+                               [(1, 0, 7), (1, 0, 7)]]:
+            a, b = np.ones(shape1), np.ones(shape2)
+            assert np.all(np.hstack((a, b)) ==
+                          np.ones((a.shape[0],
+                                   a.shape[1] + b.shape[1],
+                                   a.shape[2])))
+
     def test_dstack(self):
         import numpypy as np
-        a = np.array((1,2,3))
-        b = np.array((2,3,4))
-        c = np.dstack((a,b))
+        a = np.array((1, 2, 3))
+        b = np.array((2, 3, 4))
+        c = np.dstack((a, b))
         assert np.array_equal(c, [[[1, 2], [2, 3], [3, 4]]])
 
-        a = np.array([[1],[2],[3]])
-        b = np.array([[2],[3],[4]])
-        c = np.dstack((a,b))
+        a = np.array([[1], [2], [3]])
+        b = np.array([[2], [3], [4]])
+        c = np.dstack((a, b))
         assert np.array_equal(c, [[[1, 2]], [[2, 3]], [[3, 4]]])
+
+        for shape1, shape2 in [[(4, 2, 3), (4, 2, 7)],
+                               [(7, 2, 0), (7, 2, 10)],
+                               [(7, 2, 0), (7, 2, 0)]]:
+            a, b = np.ones(shape1), np.ones(shape2)
+            assert np.all(np.dstack((a, b)) ==
+                          np.ones((a.shape[0],
+                                   a.shape[1],
+                                   a.shape[2] + b.shape[2])))
+
+        for shape1, shape2 in [[(4, 2, 3, 5), (4, 2, 7, 5)],
+                               [(7, 2, 0, 5), (7, 2, 10, 5)],
+                               [(7, 2, 0, 5), (7, 2, 0, 5)]]:
+            a, b = np.ones(shape1), np.ones(shape2)
+            assert np.all(np.dstack((a, b)) ==
+                          np.ones((a.shape[0],
+                                   a.shape[1],
+                                   a.shape[2] + b.shape[2],
+                                   a.shape[3])))