chore: Add CompliantNamespace.is_native (#3130)

Co-authored-by: Marco Edward Gorelli <[email protected]>

Author: dangotbanned

narwhals/_compliant/namespace.py

@@ -13,7 +13,6 @@
     EagerSeriesT,
     LazyExprT,
     NativeFrameT,
-    NativeFrameT_co,
     NativeSeriesT,
 )
 from narwhals._expression_parsing import is_expr, is_series
@@ -27,7 +26,7 @@
 if TYPE_CHECKING:
     from collections.abc import Container, Iterable, Sequence
 
-    from typing_extensions import TypeAlias
+    from typing_extensions import TypeAlias, TypeIs
 
     from narwhals._compliant.selectors import CompliantSelectorNamespace
     from narwhals._compliant.when_then import CompliantWhen, EagerWhen
@@ -115,6 +114,9 @@ def selectors(self) -> CompliantSelectorNamespace[Any, Any]: ...
     def coalesce(self, *exprs: CompliantExprT) -> CompliantExprT: ...
     # NOTE: typing this accurately requires 2x more `TypeVar`s
     def from_native(self, data: Any, /) -> Any: ...
+    def is_native(self, obj: Any, /) -> TypeIs[Any]:
+        """Return `True` if `obj` can be passed to `from_native`."""
+        ...
 
 
 class DepthTrackingNamespace(
@@ -141,16 +143,18 @@ def exclude(self, excluded_names: Container[str]) -> DepthTrackingExprT:
 
 class LazyNamespace(
     CompliantNamespace[CompliantLazyFrameT, LazyExprT],
-    Protocol[CompliantLazyFrameT, LazyExprT, NativeFrameT_co],
+    Protocol[CompliantLazyFrameT, LazyExprT, NativeFrameT],
 ):
     @property
     def _backend_version(self) -> tuple[int, ...]:
         return self._implementation._backend_version()
 
     @property
     def _lazyframe(self) -> type[CompliantLazyFrameT]: ...
+    def is_native(self, obj: Any, /) -> TypeIs[NativeFrameT]:
+        return self._lazyframe._is_native(obj)
 
-    def from_native(self, data: NativeFrameT_co | Any, /) -> CompliantLazyFrameT:
+    def from_native(self, data: NativeFrameT | Any, /) -> CompliantLazyFrameT:
         if self._lazyframe._is_native(data):
             return self._lazyframe.from_native(data, context=self)
         msg = f"Unsupported type: {type(data).__name__!r}"  # pragma: no cover
@@ -173,6 +177,9 @@ def when(
         self, predicate: EagerExprT
     ) -> EagerWhen[EagerDataFrameT, EagerSeriesT, EagerExprT, NativeSeriesT]: ...
 
+    def is_native(self, obj: Any, /) -> TypeIs[NativeFrameT | NativeSeriesT]:
+        return self._dataframe._is_native(obj) or self._series._is_native(obj)
+
     @overload
     def from_native(self, data: NativeFrameT, /) -> EagerDataFrameT: ...
     @overload

narwhals/_polars/namespace.py

@@ -17,6 +17,8 @@
     from collections.abc import Iterable, Sequence
     from datetime import timezone
 
+    from typing_extensions import TypeIs
+
     from narwhals._compliant import CompliantSelectorNamespace, CompliantWhen
     from narwhals._polars.dataframe import Method, PolarsDataFrame, PolarsLazyFrame
     from narwhals._polars.typing import FrameT
@@ -101,6 +103,9 @@ def parse_into_expr(
             return self.col(data)
         return self.lit(data.to_native() if is_series(data) else data, None)
 
+    def is_native(self, obj: Any) -> TypeIs[pl.DataFrame | pl.LazyFrame | pl.Series]:
+        return isinstance(obj, (pl.DataFrame, pl.LazyFrame, pl.Series))
+
     @overload
     def from_native(self, data: pl.DataFrame, /) -> PolarsDataFrame: ...
     @overload

narwhals/_sql/namespace.py

@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING, Any, Protocol
 
 from narwhals._compliant import LazyNamespace
-from narwhals._compliant.typing import NativeExprT, NativeFrameT_co
+from narwhals._compliant.typing import NativeExprT, NativeFrameT
 from narwhals._sql.typing import SQLExprT, SQLLazyFrameT
 
 if TYPE_CHECKING:
@@ -15,8 +15,8 @@
 
 
 class SQLNamespace(
-    LazyNamespace[SQLLazyFrameT, SQLExprT, NativeFrameT_co],
-    Protocol[SQLLazyFrameT, SQLExprT, NativeFrameT_co, NativeExprT],
+    LazyNamespace[SQLLazyFrameT, SQLExprT, NativeFrameT],
+    Protocol[SQLLazyFrameT, SQLExprT, NativeFrameT, NativeExprT],
 ):
     def _function(self, name: str, *args: NativeExprT | PythonLiteral) -> NativeExprT: ...
     def _lit(self, value: Any) -> NativeExprT: ...

tests/namespace_test.py

@@ -12,7 +12,7 @@
 from narwhals._utils import Version
 
 if TYPE_CHECKING:
-    from typing_extensions import TypeAlias, assert_type
+    from typing_extensions import Never, TypeAlias, assert_type  # noqa: F401
 
     from narwhals._arrow.namespace import ArrowNamespace  # noqa: F401
     from narwhals._compliant import CompliantNamespace
@@ -177,3 +177,64 @@ class NamespaceNoVersion(Namespace): ...  # type: ignore[call-arg, type-arg]
     with pytest.raises(TypeError, match=re.compile(r"Expected.+Version.+but got.+str")):
 
         class NamespaceBadVersion(Namespace, version="invalid version"): ...  # type: ignore[arg-type, type-arg]
+
+
+def test_namespace_is_native() -> None:
+    pytest.importorskip("polars")
+    import polars as pl
+
+    unrelated: list[int] = [1, 2, 3]
+    native_1 = pl.Series(unrelated)
+    native_2 = pl.DataFrame({"a": unrelated})
+
+    maybe_native: list[pl.Series | list[int]] = [native_1, unrelated]
+    always_native = list["pl.DataFrame | pl.Series"]((native_2, native_1))
+    never_native = [unrelated, 50]
+
+    expected_maybe = [True, False]
+    expected_always = [True, True]
+    expected_never = [False, False]
+
+    ns = Namespace.from_backend("polars").compliant
+    assert [ns.is_native(el) for el in maybe_native] == expected_maybe
+    assert [ns.is_native(el) for el in always_native] == expected_always
+    assert [ns.is_native(el) for el in never_native] == expected_never
+
+    if TYPE_CHECKING:
+        if ns.is_native(native_1):
+            assert_type(native_1, "pl.Series")
+        if not ns.is_native(native_1):
+            assert_type(native_1, "Never")
+
+        if ns.is_native(unrelated):
+            # NOTE: We can't spell intersections *yet* (https://github.com/python/typing/issues/213)
+            # Would be:
+            # `<subclass of list[int] and DataFrame> | <subclass of list[int] and LazyFrame> | <subclass of list[int] and Series>``
+            assert_type(unrelated, "Never")  # pyright: ignore[reportAssertTypeFailure]
+        else:
+            assert_type(unrelated, "list[int]")
+
+        maybe_item = maybe_native[0]
+        assert_type(maybe_item, "pl.Series | list[int]")
+        if ns.is_native(maybe_item):
+            assert_type(maybe_item, "pl.Series")
+        else:
+            assert_type(maybe_item, "list[int]")
+
+        if ns.is_native(native_2):
+            assert_type(native_2, "pl.DataFrame")
+        else:
+            assert_type(native_2, "Never")
+
+        always_item = always_native[1]
+        assert_type(always_item, "pl.DataFrame | pl.Series")
+        if ns.is_native(always_item):
+            assert_type(always_item, "pl.DataFrame | pl.Series")
+            if ns._dataframe._is_native(always_item):
+                assert_type(always_item, "pl.DataFrame")
+            elif ns._series._is_native(always_item):
+                assert_type(always_item, "pl.Series")
+            else:
+                assert_type(always_item, "Never")
+        else:
+            assert_type(always_item, "Never")