Merge pull request #2897 from ContinuumIO/ndindex_fix

Fix for #2895 ndindex failing

2 files changed, 34 insertions, 2 deletions

diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py
index b07fde27d..852c5f6fd 100644
--- a/numpy/lib/index_tricks.py
+++ b/numpy/lib/index_tricks.py
@@ -502,7 +502,6 @@ class ndenumerate(object):
     def __iter__(self):
         return self
 
-
 class ndindex(object):
     """
     An N-dimensional iterator object to index arrays.
@@ -532,13 +531,36 @@ class ndindex(object):
     (2, 1, 0)
 
     """
+    # This is a hack to handle 0-d arrays correctly.
+    # 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):
+            class zero_dim_iter(object):
+                def __init__(self):
+                    self._N = 1
+                def __iter__(self):
+                    return self
+                def ndincr(self):
+                    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')
 
     def __iter__(self):
         return self
-
+        
     def ndincr(self):
         """
         Increment the multi-dimensional index by one.
diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py
index 0ede40d5a..43160ffb7 100644
--- a/numpy/lib/tests/test_index_tricks.py
+++ b/numpy/lib/tests/test_index_tricks.py
@@ -242,6 +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()