From ce5506f73bbb0233e825efef4f94dd39de05be07 Mon Sep 17 00:00:00 2001 From: "Travis E. Oliphant" Date: Mon, 21 Jan 2013 20:40:59 -0600 Subject: [PATCH 1/2] Fix-up logic for checking for zero-d arrays. --- numpy/lib/index_tricks.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py index 852c5f6fd368..15a1a559d750 100644 --- a/numpy/lib/index_tricks.py +++ b/numpy/lib/index_tricks.py @@ -535,7 +535,9 @@ class ndindex(object): # Fixing nditer would be more work but should be done eventually, # and then this entire __new__ method can be removed. def __new__(cls, *shape): - if len(shape) == 0 or (len(shape) == 1 and len(shape[0]) == 0): + if len(shape) == 1 and isinstance(shape[0], tuple): + shape = shape[0] + if len(shape) == 0: class zero_dim_iter(object): def __init__(self): self._N = 1 From 4c489f6d6edccc4d7fe2310b0e0902e980b5f52b Mon Sep 17 00:00:00 2001 From: "Travis E. Oliphant" Date: Mon, 21 Jan 2013 20:50:15 -0600 Subject: [PATCH 2/2] TST: Add a test for ndindex call. --- numpy/lib/tests/test_index_tricks.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py index 43160ffb7727..a6e65ef56498 100644 --- a/numpy/lib/tests/test_index_tricks.py +++ b/numpy/lib/tests/test_index_tricks.py @@ -245,6 +245,10 @@ def test_ndindex(): x = list(np.ndindex((1, 2, 3))) assert_array_equal(x, expected) + # Test use of scalars and tuples + x = list(np.ndindex((3,))) + assert_array_equal(x, list(np.ndindex(3))) + # Make sure size argument is optional x = list(np.ndindex()) assert_equal(x, [()])