diff options
| author | Dominic Jack <thedomjack@gmail.com> | 2019-01-30 19:36:50 +1000 |
|---|---|---|
| committer | Matti Picus <matti.picus@gmail.com> | 2019-01-30 11:36:50 +0200 |
| commit | 1be850b0782a1dcbb3da22fb9d5c74c3ce367936 (patch) | |
| tree | 0a6370ef1ac3bd94d715f29557f6bcf74c1b1d7b | |
| parent | 34728c4efe2e52dcad7d675b7cbf881a735db7ec (diff) | |
| download | numpy-1be850b0782a1dcbb3da22fb9d5c74c3ce367936.tar.gz | |
BUG: squeeze corner case
| -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]] |