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

Code before:

    x = list(ndindex((1, 2, 3)))
    assert_array_equal(x, expected)

    # Test use of scalars and tuples
    x = list(ndindex((3,)))
    assert_array_equal(x, list(ndindex(3)))

    # 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.  Empty Shapes and Zero Dimensions
def test_ndindex_empty_and_zero_dimensions():
    # Empty shape
    assert list(np.ndindex()) == [()]
    assert list(np.ndindex(())) == [()]

    # Single dimension
    assert list(np.ndindex(4)) == [(0,), (1,), (2,), (3,)]

    # Zero dimensions should produce empty iterator
    assert list(np.ndindex(0, 3)) == []
    assert list(np.ndindex(3, 0, 2)) == []

# 2. Non-Integer Dimensions Raise TypeError
def test_ndindex_non_integer_dimensions():
    """Test that non-integer dimensions raise TypeError."""
    with pytest.raises(TypeError):
        list(np.ndindex(2.5))
    with pytest.raises(TypeError):
        list(np.ndindex("2"))
    with pytest.raises(TypeError):
        list(np.ndindex([2, 3]))
    with pytest.raises(TypeError):
        list(np.ndindex((2.0, 3)))

# 3. StopIteration Behavior
def test_ndindex_stop_iteration_behavior():
    """Test that StopIteration is raised properly after exhaustion."""
    it = np.ndindex(2, 2)

Review comment (suggestion):
It probably does make sense to add a few more test cases, but note that the already-existing `test_ndindex` already has this particular test case for example, and may cover one or two other cases.

Code after:

    x = list(ndindex((1, 2, 3)))
    assert_array_equal(x, expected)

    # Test use of scalars and tuples
    x = list(ndindex((3,)))
    assert_array_equal(x, list(ndindex(3)))

    # 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)