Repository: scikit-learn/scikit-learn
File: sklearn/utils/tests/test_array_api.py (Python)

Code before:
from sklearn.utils.fixes import _IS_32BIT, CSR_CONTAINERS, np_version, parse_version


@pytest.mark.parametrize("X", [numpy.asarray([1, 2, 3]), [1, 2, 3], (1, 2, 3)])
def test_get_namespace_ndarray_or_similar_default(X):
    """Check that get_namespace returns NumPy wrapper"""
    xp_out, is_array_api_compliant = get_namespace(X)
    assert xp_out is np_compat
    assert not is_array_api_compliant


def test_get_namespace_ndarray_creation_device():
    """Check expected behavior with device and creation functions."""
    X = numpy.asarray([1, 2, 3])
    xp_out, _ = get_namespace(X)

    full_array = xp_out.full(10, fill_value=2.0, device="cpu")
    assert_allclose(full_array, [2.0] * 10)

    with pytest.raises(ValueError, match="Unsupported device"):
        xp_out.zeros(10, device="cuda")


@skip_if_array_api_compat_not_configured
@pytest.mark.parametrize("X", [numpy.asarray([1, 2, 3]), [1, 2, 3], (1, 2, 3)])
def test_get_namespace_ndarray_with_dispatch(X):
    """Test get_namespace on NumPy ndarrays."""

    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(X)
        assert is_array_api_compliant == isinstance(X, numpy.ndarray)

        # In the future, NumPy should become API compliant library and we should have
        # assert xp_out is numpy
        assert xp_out is np_compat


@skip_if_array_api_compat_not_configured
@pytest.mark.parametrize(
    "constructor_name", ["pyarrow", "dataframe", "polars", "series"]
)
def test_get_namespace_df_with_dispatch(constructor_name):
    """Test get_namespace on dataframes and series."""

    df = _convert_container([[1, 4, 2], [3, 3, 6]], constructor_name)
    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(df)
        assert not is_array_api_compliant

        # When operating on dataframes or series the Numpy namespace is
        # the right thing to use.

Review comment (suggestion):
nit, to match test above
```suggestion
def test_get_namespace_ndarray_or_similar_default_with_dispatch(X):
```

Code after:
from sklearn.utils.fixes import _IS_32BIT, CSR_CONTAINERS, np_version, parse_version


@pytest.mark.parametrize("X", [numpy.asarray([1, 2, 3]), [1, 2, 3], (1, 2, 3)])
def test_get_namespace_ndarray_or_similar_default(X):
    """Check that get_namespace returns NumPy wrapper"""
    xp_out, is_array_api_compliant = get_namespace(X)
    assert xp_out is np_compat
    assert not is_array_api_compliant


def test_get_namespace_ndarray_creation_device():
    """Check expected behavior with device and creation functions."""
    X = numpy.asarray([1, 2, 3])
    xp_out, _ = get_namespace(X)

    full_array = xp_out.full(10, fill_value=2.0, device="cpu")
    assert_allclose(full_array, [2.0] * 10)

    with pytest.raises(ValueError, match="Unsupported device"):
        xp_out.zeros(10, device="cuda")


@skip_if_array_api_compat_not_configured
@pytest.mark.parametrize("X", [numpy.asarray([1, 2, 3]), [1, 2, 3], (1, 2, 3)])
def test_get_namespace_ndarray_or_similar_default_with_dispatch(X):
    """Test get_namespace on NumPy ndarrays."""

    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(X)
        assert is_array_api_compliant == isinstance(X, numpy.ndarray)

        # In the future, NumPy should become API compliant library and we should have
        # assert xp_out is numpy
        assert xp_out is np_compat


@skip_if_array_api_compat_not_configured
@pytest.mark.parametrize(
    "constructor_name", ["pyarrow", "dataframe", "polars", "series"]
)
def test_get_namespace_df_with_dispatch(constructor_name):
    """Test get_namespace on dataframes and series."""

    df = _convert_container([[1, 4, 2], [3, 3, 6]], constructor_name)
    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(df)
        assert not is_array_api_compliant

        # When operating on dataframes or series the Numpy namespace is
        # the right thing to use.