Repository: narwhals-dev/narwhals
Pull request #1520 (2024-12-06 11:47:40)

 enh: get_native_namespace accepts native objects
closes https://github.com/narwhals-dev/narwhals/issues/1498

making this PR even though i'd marked it "reserved for sprint" because I need it for something else i'm working on

<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues

- Related issue #\<issue number\>
- Closes #\<issue number\>

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below



Patch:
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index 60aa7c7d..aacc2a2f 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -5,7 +5,6 @@ from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterable
 from typing import Literal
-from typing import cast
 
 from narwhals._arrow.dataframe import ArrowDataFrame
 from narwhals._arrow.expr import ArrowExpr
@@ -447,18 +446,17 @@ class ArrowWhen:
 
         plx = ArrowNamespace(backend_version=self._backend_version, version=self._version)
 
-        condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0]  # type: ignore[arg-type]
+        condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0]
         try:
-            value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0]  # type: ignore[arg-type]
+            value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0]
         except TypeError:
             # `self._otherwise_value` is a scalar and can't be converted to an expression
-            value_series = condition.__class__._from_iterable(  # type: ignore[call-arg]
+            value_series = condition.__class__._from_iterable(
                 [self._then_value] * len(condition),
                 name="literal",
                 backend_version=self._backend_version,
                 version=self._version,
             )
-        value_series = cast(ArrowSeries, value_series)
 
         value_series_native = value_series._native_series
         condition_native = condition._native_series.combine_chunks()
@@ -475,7 +473,7 @@ class ArrowWhen:
         try:
             otherwise_series = parse_into_expr(
                 self._otherwise_value, namespace=plx
-            )._call(df)[0]  # type: ignore[arg-type]
+            )._call(df)[0]
         except TypeError:
             # `self._otherwise_value` is a scalar and can't be converted to an expression.
             # Remark that string values _are_ converted into expressions!
@@ -487,8 +485,6 @@ class ArrowWhen:
                 )
             ]
         else:
-            otherwise_series = cast(ArrowSeries, otherwise_series)
-            condition = cast(ArrowSeries, condition)
             condition_native, otherwise_native = broadcast_series(
                 [condition, otherwise_series]
             )
diff --git a/narwhals/_dask/namespace.py b/narwhals/_dask/namespace.py
index 3d5fc51b..3478a0a4 100644
--- a/narwhals/_dask/namespace.py
+++ b/narwhals/_dask/namespace.py
@@ -436,10 +436,10 @@ class DaskWhen:
 
         plx = DaskNamespace(backend_version=self._backend_version, version=self._version)
 
-        condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0]  # type: ignore[arg-type]
+        condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0]
         condition = cast("dask_expr.Series", condition)
         try:
-            value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0]  # type: ignore[arg-type]
+            value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0]
         except TypeError:
             # `self._otherwise_value` is a scalar and can't be converted to an expression
             _df = condition.to_frame("a")
@@ -453,7 +453,7 @@ class DaskWhen:
         try:
             otherwise_series = parse_into_expr(
                 self._otherwise_value, namespace=plx
-            )._call(df)[0]  # type: ignore[arg-type]
+            )._call(df)[0]
         except TypeError:
             # `self._otherwise_value` is a scalar and can't be converted to an expression
             return [value_series.where(condition, self._otherwise_value)]
diff --git a/narwhals/_expression_parsing.py b/narwhals/_expression_parsing.py
index e10a4730..9321641b 100644
--- a/narwhals/_expression_parsing.py
+++ b/narwhals/_expression_parsing.py
@@ -75,7 +75,7 @@ def evaluate_into_expr(
     df: CompliantDataFrame, into_expr: IntoCompliantExpr
 ) -> ListOfCompliantSeries:
     """Return list of raw columns."""
-    expr = parse_into_expr(into_expr, namespace=df.__narwhals_namespace__())
+    expr = parse_into_expr(into_expr, namespace=df.__narwhals_namespace__())  # type: ignore[arg-type]
     return expr._call(df)  # type: ignore[arg-type]
 
 
@@ -183,13 +183,31 @@ def parse_into_exprs(
 
     See `parse_into_expr` for more details.
     """
-    return [parse_into_expr(into_expr, namespace=namespace) for into_expr in exprs] + [
-        parse_into_expr(expr, namespace=namespace).alias(name)
+    return [parse_into_expr(into_expr, namespace=namespace) for into_expr in exprs] + [  # type: ignore[arg-type]
+        parse_into_expr(expr, namespace=namespace).alias(name)  # type: ignore[arg-type]
         for name, expr in named_exprs.items()
     ]
 
 
+@overload
+def parse_into_expr(into_expr: IntoArrowExpr, namespace: ArrowNamespace) -> ArrowExpr: ...
+@overload
 def parse_into_expr(
+    into_expr: IntoPandasLikeExpr, namespace: PandasLikeNamespace
+) -> PandasLikeExpr: ...
+@overload
+def parse_into_expr(
+    into_expr: IntoPolarsExpr, namespace: PolarsNamespace
+) -> PolarsExpr: ...
+@overload
+def parse_into_expr(
+    into_expr: IntoSparkLikeExpr, namespace: SparkLikeNamespace
+) -> SparkLikeExpr: ...
+@overload
+def parse_into_expr(into_expr: IntoDaskExpr, namespace: DaskNamespace) -> DaskExpr: ...
+
+
+def parse_into_expr(  # type: ignore[misc]
     into_expr: IntoCompliantExpr,
     *,
     namespace: CompliantNamespace,
diff --git a/narwhals/_pandas_like/namespace.py b/narwhals/_pandas_like/namespace.py
index 880d4232..817bc4a9 100644
--- a/narwhals/_pandas_like/namespace.py
+++ b/narwhals/_pandas_like/namespace.py
@@ -6,7 +6,6 @@ from typing import Any
 from typing import Callable
 from typing import Iterable
 from typing import Literal
-from typing import cast
 
 from narwhals._expression_parsing import combine_root_names
 from narwhals._expression_parsing import parse_into_exprs
@@ -486,12 +485,12 @@ class PandasWhen:
             version=self._version,
         )
 
-        condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0]  # type: ignore[arg-type]
+        condition = parse_into_expr(self._condition, namespace=plx)._call(df)[0]
         try:
-            value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0]  # type: ignore[arg-type]
+            value_series = parse_into_expr(self._then_value, namespace=plx)._call(df)[0]
         except TypeError:
             # `self._otherwise_value` is a scalar and can't be converted to an expression
-            value_series = condition.__class__._from_iterable(  # type: ignore[call-arg]
+            value_series = condition.__class__._from_iterable(
                 [self._then_value] * len(condition),
                 name="literal",
                 index=condition._native_series.index,
@@ -499,8 +498,6 @@ class PandasWhen:
                 backend_version=self._backend_version,
                 version=self._version,
             )
-        value_series = cast(PandasLikeSeries, value_series)
-
         value_series_native, condition_native = broadcast_align_and_extract_native(
             value_series, condition
         )
@@ -514,7 +511,7 @@ class PandasWhen:
         try:
             otherwise_series = parse_into_expr(
                 self._otherwise_value, namespace=plx
-            )._call(df)[0]  # type: ignore[arg-type]
+            )._call(df)[0]
         except TypeError:
             # `self._otherwise_value` is a scalar and can't be converted to an expression
             return [
diff --git a/narwhals/translate.py b/narwhals/translate.py
index 2bd0949f..e673d7e3 100644
--- a/narwhals/translate.py
+++ b/narwhals/translate.py
@@ -38,6 +38,10 @@ from narwhals.dependencies import is_pyspark_dataframe
 from narwhals.utils import Version
 
 if TYPE_CHECKING:
+    import pandas as pd
+    import polars as pl
+    import pyarrow as pa
+
     from narwhals.dataframe import DataFrame
     from narwhals.dataframe import LazyFrame
     from narwhals.series import Series
@@ -749,7 +753,18 @@ def _from_native_impl(  # noqa: PLR0915
     return native_object
 
 
-def get_native_namespace(obj: DataFrame[Any] | LazyFrame[Any] | Series[Any]) -> Any:
+def get_native_namespace(
+    obj: DataFrame[Any]
+    | LazyFrame[Any]
+    | Series[Any]
+    | pd.DataFrame
+    | pd.Series
+    | pl.DataFrame
+    | pl.LazyFrame
+    | pl.Series
+    | pa.Table
+    | pa.ChunkedArray,
+) -> Any:
     """Get native namespace from object.
 
     Arguments:
@@ -769,7 +784,22 @@ def get_native_namespace(obj: DataFrame[Any] | LazyFrame[Any] | Series[Any]) ->
         >>> nw.get_native_namespace(df)
         <module 'polars'...>
     """
-    return obj.__native_namespace__()
+    if hasattr(obj, "__native_namespace__"):
+        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)
 
 
 def narwhalify(


Test patch:
diff --git a/tests/translate/get_native_namespace_test.py b/tests/translate/get_native_namespace_test.py
index f02c4c8d..5ad3f94e 100644
--- a/tests/translate/get_native_namespace_test.py
+++ b/tests/translate/get_native_namespace_test.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import pandas as pd
 import polars as pl
 import pyarrow as pa
+import pytest
 
 import narwhals.stable.v1 as nw
 
@@ -10,7 +11,19 @@ import narwhals.stable.v1 as nw
 def test_native_namespace() -> None:
     df = 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
     df = nw.from_native(pd.DataFrame({"a": [1, 2, 3]}))
     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
     df = nw.from_native(pa.table({"a": [1, 2, 3]}))
     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
+
+
+def test_get_native_namespace_invalid() -> None:
+    with pytest.raises(TypeError, match="Could not get native namespace"):
+        nw.get_native_namespace(1)