diff options
Diffstat (limited to 'numpy/lib/tests/test_index_tricks.py')
| -rw-r--r-- | numpy/lib/tests/test_index_tricks.py | 8 |
1 files changed, 7 insertions, 1 deletions
diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py index beda2d146..0ede40d5a 100644 --- a/numpy/lib/tests/test_index_tricks.py +++ b/numpy/lib/tests/test_index_tricks.py @@ -2,7 +2,7 @@ from numpy.testing import * import numpy as np from numpy import ( array, ones, r_, mgrid, unravel_index, zeros, where, ndenumerate, fill_diagonal, diag_indices, - diag_indices_from, s_, index_exp ) + diag_indices_from, s_, index_exp, ndindex ) class TestRavelUnravelIndex(TestCase): def test_basic(self): @@ -237,5 +237,11 @@ def test_diag_indices_from(): assert_array_equal(c, np.arange(4)) +def test_ndindex(): + x = list(np.ndindex(1, 2, 3)) + expected = [ix for ix, e in np.ndenumerate(np.zeros((1, 2, 3)))] + assert_array_equal(x, expected) + + if __name__ == "__main__": run_module_suite() |