fix: Include all backends in get_native_namespace (#2608)

* fix: Include all backends in get_native_namespace

* test typing and pyarrow_table specification

* no cover cudf and None return in test

Author: FBruzzesi

narwhals/translate.py

@@ -12,30 +12,18 @@
     is_native_spark_like,
 )
 from narwhals.dependencies import (
-    get_cudf,
     get_dask,
     get_dask_expr,
-    get_modin,
     get_numpy,
     get_pandas,
-    get_polars,
-    get_pyarrow,
-    is_cudf_dataframe,
-    is_cudf_series,
     is_cupy_scalar,
     is_dask_dataframe,
     is_duckdb_relation,
     is_ibis_table,
-    is_modin_dataframe,
-    is_modin_series,
     is_numpy_scalar,
-    is_pandas_dataframe,
     is_pandas_like_dataframe,
-    is_pandas_series,
-    is_polars_dataframe,
     is_polars_lazyframe,
     is_polars_series,
-    is_pyarrow_chunked_array,
     is_pyarrow_scalar,
     is_pyarrow_table,
 )
@@ -602,25 +590,20 @@ def get_native_namespace(
     return result.pop()
 
 
-def _get_native_namespace_single_obj(  # noqa: PLR0911
+def _get_native_namespace_single_obj(
     obj: DataFrame[Any] | LazyFrame[Any] | Series[Any] | IntoFrame | IntoSeries,
 ) -> Any:
+    from contextlib import suppress
+
     from narwhals.utils import has_native_namespace
 
+    with suppress(TypeError, AssertionError):
+        return Version.MAIN.namespace.from_native_object(
+            obj
+        ).implementation.to_native_namespace()
+
     if has_native_namespace(obj):
         return obj.__native_namespace__()
-    if is_pandas_dataframe(obj) or is_pandas_series(obj):
-        return get_pandas()
-    if is_modin_dataframe(obj) or is_modin_series(obj):  # pragma: no cover
-        return get_modin()
-    if is_pyarrow_table(obj) or is_pyarrow_chunked_array(obj):
-        return get_pyarrow()
-    if is_cudf_dataframe(obj) or is_cudf_series(obj):  # pragma: no cover
-        return get_cudf()
-    if is_dask_dataframe(obj):  # pragma: no cover
-        return get_dask()
-    if is_polars_dataframe(obj) or is_polars_lazyframe(obj) or is_polars_series(obj):
-        return get_polars()
     msg = f"Could not get native namespace from object of type: {type(obj)}"
     raise TypeError(msg)
 

tests/translate/get_native_namespace_test.py

@@ -1,46 +1,82 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import pytest
 
 import narwhals as nw
 
 if TYPE_CHECKING:
     from narwhals.typing import Frame
+    from tests.utils import Constructor
 
 
-def test_native_namespace_polars() -> None:
-    pytest.importorskip("polars")
-    import polars as pl
+data = {"a": [1, 2, 3]}
 
-    df: Frame = nw.from_native(pl.DataFrame({"a": [1, 2, 3]}))
-    assert nw.get_native_namespace(df) is pl
-    assert nw.get_native_namespace(df.to_native()) is pl
-    assert nw.get_native_namespace(df.lazy().to_native()) is pl
-    assert nw.get_native_namespace(df["a"].to_native()) is pl
-    assert nw.get_native_namespace(df, df["a"].to_native()) is pl
 
+def _get_expected_namespace(constructor_name: str) -> Any | None:  # noqa: PLR0911
+    """Get expected namespace module for a given constructor."""
+    if "pandas" in constructor_name:
+        import pandas as pd
+
+        return pd
+    elif "polars" in constructor_name:
+        import polars as pl
+
+        return pl
+    elif "pyarrow_table" in constructor_name:
+        import pyarrow as pa
+
+        return pa
+    elif "duckdb" in constructor_name:
+        import duckdb
+
+        return duckdb
+    elif "cudf" in constructor_name:  # pragma: no cover
+        import cudf
+
+        return cudf
+    elif "modin" in constructor_name:
+        import modin.pandas as mpd
+
+        return mpd
+    elif "dask" in constructor_name:
+        import dask.dataframe as dd
+
+        return dd
+    elif "ibis" in constructor_name:
+        import ibis
+
+        return ibis
+    elif "sqlframe" in constructor_name:
+        import sqlframe
+
+        return sqlframe
+    return None  # pragma: no cover
+
+
+def test_native_namespace_frame(constructor: Constructor) -> None:
+    constructor_name = constructor.__name__
+    if constructor_name == "pyspark_lazy_constructor":
+        pytest.skip(reason="Requires special handling for spark local vs spark connect")
+
+    expected_namespace = _get_expected_namespace(constructor_name=constructor_name)
+
+    df: Frame = nw.from_native(constructor(data))
+    assert nw.get_native_namespace(df) is expected_namespace
+    assert nw.get_native_namespace(df.to_native()) is expected_namespace
+    assert nw.get_native_namespace(df.lazy().to_native()) is expected_namespace
 
-def test_native_namespace_pandas() -> None:
-    pytest.importorskip("pandas")
-    import pandas as pd
 
-    df: Frame = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}), eager_only=True)
-    assert nw.get_native_namespace(df) is pd
-    assert nw.get_native_namespace(df.to_native()) is pd
-    assert nw.get_native_namespace(df["a"].to_native()) is pd
-    assert nw.get_native_namespace(df, df["a"].to_native()) is pd
+def test_native_namespace_series(constructor_eager: Constructor) -> None:
+    constructor_name = constructor_eager.__name__
 
+    expected_namespace = _get_expected_namespace(constructor_name=constructor_name)
 
-def test_native_namespace_pyarrow() -> None:
-    pytest.importorskip("pyarrow")
-    import pyarrow as pa
+    df: Frame = nw.from_native(constructor_eager(data), eager_only=True)
 
-    df: Frame = nw.from_native(pa.table({"a": [1, 2, 3]}), eager_only=True)
-    assert nw.get_native_namespace(df) is pa
-    assert nw.get_native_namespace(df.to_native()) is pa
-    assert nw.get_native_namespace(df, df["a"].to_native()) is pa
+    assert nw.get_native_namespace(df["a"].to_native()) is expected_namespace
+    assert nw.get_native_namespace(df, df["a"].to_native()) is expected_namespace
 
 
 def test_get_native_namespace_invalid() -> None: