Fix the test for numpy.ndindex()

Fix ndindex for 0-d arrays.
Add tests for tuple arguments to ndindex

2 files changed, 24 insertions, 5 deletions

diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py
index 139883847..6bb06b7b6 100644
--- a/numpy/lib/index_tricks.py
+++ b/numpy/lib/index_tricks.py
@@ -532,16 +532,29 @@ class ndindex(object):
 
     """
     # This is a hack to handle 0-d arrays correctly.
-    # Fixing nditer would be more work but should be done eventually.
+    # 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:
-            def zerodim_gen():
-                yield ()
-            return zerodim_gen()
+        if len(shape) == 0 or (len(shape) == 1 and len(shape[0]) == 0):
+            class zero_dim_iter(object):
+                def __init__(self):
+                    self._N = 1
+                def __iter__(self):
+                    return self
+                def ndincr(self):
+                    return self.next()
+                def next(self):
+                    if self._N > 0:
+                        self._N -= 1
+                        return ()
+                    raise StopIteration
+            return zero_dim_iter()
         else:
             return super(ndindex, cls).__new__(cls)
             
     def __init__(self, *shape):
+        if len(shape) == 1 and isinstance(shape[0], tuple):
+            shape = shape[0]
         x = as_strided(_nx.zeros(1), shape=shape, strides=_nx.zeros_like(shape))
         self._it = _nx.nditer(x, flags=['multi_index'], order='C')
 
diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py
index 3cb06b4fb..43160ffb7 100644
--- a/numpy/lib/tests/test_index_tricks.py
+++ b/numpy/lib/tests/test_index_tricks.py
@@ -242,10 +242,16 @@ def test_ndindex():
     expected = [ix for ix, e in np.ndenumerate(np.zeros((1, 2, 3)))]
     assert_array_equal(x, expected)
 
+    x = list(np.ndindex((1, 2, 3)))
+    assert_array_equal(x, expected)
+
     # Make sure size argument is optional
     x = list(np.ndindex())
     assert_equal(x, [()])
 
+    x = list(np.ndindex(()))
+    assert_equal(x, [()])    
+
 
 if __name__ == "__main__":
     run_module_suite()