feat: support passing multiple arguments positionally to get_native_namespace (#2178)

Author: MarcoGorelli

narwhals/translate.py

@@ -40,11 +40,6 @@
 from narwhals.utils import Version
 
 if TYPE_CHECKING:
-    import pandas as pd
-    import polars as pl
-    import pyarrow as pa
-
-    from narwhals._arrow.typing import ArrowChunkedArray
     from narwhals.dataframe import DataFrame
     from narwhals.dataframe import LazyFrame
     from narwhals.series import Series
@@ -790,21 +785,14 @@ def _from_native_impl(  # noqa: PLR0915
 
 
 def get_native_namespace(
-    obj: DataFrame[Any]
-    | LazyFrame[Any]
-    | Series[Any]
-    | pd.DataFrame
-    | pd.Series[Any]
-    | pl.DataFrame
-    | pl.LazyFrame
-    | pl.Series
-    | pa.Table
-    | ArrowChunkedArray,
+    *obj: DataFrame[Any] | LazyFrame[Any] | Series[Any] | IntoFrame | IntoSeries,
 ) -> Any:
     """Get native namespace from object.
 
     Arguments:
-        obj: Dataframe, Lazyframe, or Series.
+        obj: Dataframe, Lazyframe, or Series. Multiple objects can be
+            passed positionally, in which case they must all have the
+            same native namespace (else an error is raised).
 
     Returns:
         Native module.
@@ -820,6 +808,19 @@ def get_native_namespace(
         >>> nw.get_native_namespace(df)
         <module 'polars'...>
     """
+    if not obj:
+        msg = "At least one object must be passed to `get_native_namespace`."
+        raise ValueError(msg)
+    result = {_get_native_namespace_single_obj(x) for x in obj}
+    if len(result) != 1:
+        msg = f"Found objects with different native namespaces: {result}."
+        raise ValueError(msg)
+    return result.pop()
+
+
+def _get_native_namespace_single_obj(
+    obj: DataFrame[Any] | LazyFrame[Any] | Series[Any] | IntoFrame | IntoSeries,
+) -> Any:
     from narwhals.utils import has_native_namespace
 
     if has_native_namespace(obj):

tests/translate/get_native_namespace_test.py

@@ -19,16 +19,22 @@ def test_native_namespace() -> None:
     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
     df = 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
     df = 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["a"].to_native()) is pa
+    assert nw.get_native_namespace(df, df["a"].to_native()) is pa
 
 
 def test_get_native_namespace_invalid() -> None:
     with pytest.raises(TypeError, match="Could not get native namespace"):
         nw.get_native_namespace(1)  # type: ignore[arg-type]
+    with pytest.raises(ValueError, match="At least one object"):
+        nw.get_native_namespace()
+    with pytest.raises(ValueError, match="Found objects with different"):
+        nw.get_native_namespace(pd.Series([1]), pl.Series([2]))