Skip to content

Commit

Permalink
Merge pull request numpy#2940 from ContinuumIO/ndindex_fix_more
Browse files Browse the repository at this point in the history
BUG: Fix logic in ndindex to match __init__ method.   Fixes bug in ndindex((3,))
  • Loading branch information
teoliphant committed Jan 24, 2013
2 parents dce1018 + 4c489f6 commit 2b4c9a6
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 1 deletion.
4 changes: 3 additions & 1 deletion numpy/lib/index_tricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions numpy/lib/tests/test_index_tricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, [()])
Expand Down

0 comments on commit 2b4c9a6

Please sign in to comment.