Fix the test for numpy.ndindex()

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

1 files changed, 18 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')