diff options
Diffstat (limited to 'numpy')
| -rw-r--r-- | numpy/core/fromnumeric.py | 2 | ||||
| -rw-r--r-- | numpy/core/tests/test_numeric.py | 10 |
2 files changed, 10 insertions, 2 deletions
diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py index d94372986..7b179cf31 100644 --- a/numpy/core/fromnumeric.py +++ b/numpy/core/fromnumeric.py @@ -1385,7 +1385,7 @@ def squeeze(a, axis=None): try: squeeze = a.squeeze except AttributeError: - return _wrapit(a, 'squeeze') + return _wrapit(a, 'squeeze', axis=axis) if axis is None: return squeeze() else: diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py index 37534720a..ae11b1cab 100644 --- a/numpy/core/tests/test_numeric.py +++ b/numpy/core/tests/test_numeric.py @@ -152,7 +152,15 @@ class TestNonarrayArgs(object): def test_squeeze(self): A = [[[1, 1, 1], [2, 2, 2], [3, 3, 3]]] - assert_(np.squeeze(A).shape == (3, 3)) + assert_equal(np.squeeze(A).shape, (3, 3)) + assert_equal(np.squeeze(np.zeros((1, 3, 1))).shape, (3,)) + assert_equal(np.squeeze(np.zeros((1, 3, 1)), axis=0).shape, (3, 1)) + assert_equal(np.squeeze(np.zeros((1, 3, 1)), axis=-1).shape, (1, 3)) + assert_equal(np.squeeze(np.zeros((1, 3, 1)), axis=2).shape, (1, 3)) + assert_equal(np.squeeze([np.zeros((3, 1))]).shape, (3,)) + assert_equal(np.squeeze([np.zeros((3, 1))], axis=0).shape, (3, 1)) + assert_equal(np.squeeze([np.zeros((3, 1))], axis=2).shape, (1, 3)) + assert_equal(np.squeeze([np.zeros((3, 1))], axis=-1).shape, (1, 3)) def test_std(self): A = [[1, 2, 3], [4, 5, 6]] |