Repository: numpy/numpy
File: numpy/lib/tests/test_index_tricks.py (Python)

Code before:

    # Make sure size argument is optional
    x = list(ndindex())
    assert_equal(x, [()])

    x = list(ndindex(()))
    assert_equal(x, [()])

    # Make sure 0-sized ndindex works correctly
    x = list(ndindex(*[0]))
    assert_equal(x, [])


# ndindex test cases for reimplementation of ndindex using itertools
# 1. Zero Dimensions (explicitly zero-length dimensions)
# This test focuses specifically on cases where one or more dimensions are zero,
# ensuring the iterator produces an empty sequence. This avoids redundancy
# with basic tests handled by the main `test_ndindex` function.
def test_ndindex_zero_dimensions_explicit():
    """Test ndindex produces empty iterators for explicit zero-length dimensions."""
    assert list(np.ndindex(0, 3)) == []
    assert list(np.ndindex(3, 0, 2)) == []
    assert list(np.ndindex(0)) == []  # Test for a single zero dimension


# 2. Non-Integer Dimensions Raise TypeError
@pytest.mark.parametrize("bad_shape", [
    2.5, "2", [2, 3], (2.0, 3)
])
def test_ndindex_non_integer_dimensions(bad_shape):
    """Test that non-integer dimensions raise TypeError."""
    with pytest.raises(TypeError):
        # Passing invalid_shape_arg directly to ndindex, it will try to use it as a dimension.
        # This should trigger the TypeError during object initialization or iteration.
        list(np.ndindex(bad_shape))

# 3. StopIteration Behavior
def test_ndindex_stop_iteration_behavior():
    """Test that StopIteration is raised properly after exhaustion."""
    it = np.ndindex(2, 2)
    # Exhaust the iterator
    list(it)
    # Should raise StopIteration on subsequent calls
    with pytest.raises(StopIteration):
        next(it)

# 4. Iterator Independence
def test_ndindex_iterator_independence():
    """Test that each ndindex instance creates independent iterators."""
    shape = (2, 3)
    iter1 = np.ndindex(*shape)

Review comment (suggestion):
Please remove these comments before the tests. If additional info is needed it can be part of the doc string

```suggestion
```

Code after:

    # Make sure size argument is optional
    x = list(ndindex())
    assert_equal(x, [()])

    x = list(ndindex(()))
    assert_equal(x, [()])

    # Make sure 0-sized ndindex works correctly
    x = list(ndindex(*[0]))
    assert_equal(x, [])


def test_ndindex_zero_dimensions_explicit():
    """Test ndindex produces empty iterators for explicit
    zero-length dimensions."""
    assert list(np.ndindex(0, 3)) == []
    assert list(np.ndindex(3, 0, 2)) == []
    assert list(np.ndindex(0)) == []


@pytest.mark.parametrize("bad_shape", [2.5, "2", [2, 3], (2.0, 3)])
def test_ndindex_non_integer_dimensions(bad_shape):
    """Test that non-integer dimensions raise TypeError."""
    with pytest.raises(TypeError):
        # Passing invalid_shape_arg directly to ndindex. It will try to use it
        # as a dimension and should trigger a TypeError.
        list(np.ndindex(bad_shape))


def test_ndindex_stop_iteration_behavior():
    """Test that StopIteration is raised properly after exhaustion."""
    it = np.ndindex(2, 2)
    # Exhaust the iterator
    list(it)
    # Should raise StopIteration on subsequent calls
    with pytest.raises(StopIteration):
        next(it)


def test_ndindex_iterator_independence():
    """Test that each ndindex instance creates independent iterators."""
    shape = (2, 3)
    iter1 = np.ndindex(*shape)
    iter2 = np.ndindex(*shape)

    next(iter1)
    next(iter1)

    assert_equal(next(iter2), (0, 0))
    assert_equal(next(iter1), (0, 2))