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

Code before:
    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
def test_get_namespace_ndarray_with_dispatch():
    """Test get_namespace on NumPy ndarrays."""

    X_np = numpy.asarray([[1, 2, 3]])

    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(X_np)
        assert is_array_api_compliant

        # 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
def test_get_namespace_pandas_with_dispatch():
    """Test get_namespace on pandas dataframes and series."""
    pd = pytest.importorskip("pandas")

    X_pd = pd.DataFrame([[1, 2, 3]])
    X_pd_series = pd.Series([1, 2, 3])

    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(X_pd)
        assert not is_array_api_compliant

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

        xp_out, is_array_api_compliant = get_namespace(X_pd_series)
        assert not is_array_api_compliant
        assert xp_out is np_compat


@skip_if_array_api_compat_not_configured
def test_get_namespace_sparse_with_dispatch():
    """Test get_namespace on sparse arrays."""
    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(sp.csr_array([[1, 2, 3]]))
        assert not is_array_api_compliant

Review comment (suggestion):
Investigate parametrizing this on pandas, pyarrow, polars instead of having one test per dataframe library.

Code after:
    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
def test_get_namespace_ndarray_with_dispatch():
    """Test get_namespace on NumPy ndarrays."""

    X_np = numpy.asarray([[1, 2, 3]])

    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(X_np)
        assert is_array_api_compliant

        # 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.
        assert xp_out is np_compat


@skip_if_array_api_compat_not_configured
def test_get_namespace_sparse_with_dispatch():
    """Test get_namespace on sparse arrays."""
    with config_context(array_api_dispatch=True):
        xp_out, is_array_api_compliant = get_namespace(sp.csr_array([[1, 2, 3]]))
        assert not is_array_api_compliant

        # When operating on sparse arrays the Numpy namespace is
        # the right thing to use.