diff options
Diffstat (limited to 'numpy/lib/index_tricks.py')
| -rw-r--r-- | numpy/lib/index_tricks.py | 23 |
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') |