feat(typing): Make Implementation less opaque

narwhals/_namespace.py

@@ -10,6 +10,7 @@
     Literal,
     Protocol,
     TypeVar,
+    cast,
     overload,
 )
 
@@ -37,8 +38,6 @@
     import pandas as pd
     import polars as pl
     import pyarrow as pa
-    import pyspark.sql as pyspark_sql
-    from pyspark.sql.connect.dataframe import DataFrame as PySparkConnectDataFrame
     from typing_extensions import Self, TypeAlias, TypeIs
 
     from narwhals._arrow.namespace import ArrowNamespace
@@ -91,7 +90,7 @@ class _BasePandasLikeFrame(NativeDataFrame, _BasePandasLike, Protocol): ...
     class _BasePandasLikeSeries(NativeSeries, _BasePandasLike, Protocol):
         def where(self, cond: Any, other: Any = ..., **kwds: Any) -> Any: ...
 
-    class _NativeDask(Protocol):
+    class _NativeDask(NativeLazyFrame, Protocol):
         _partition_type: type[pd.DataFrame]
 
     class _CuDFDataFrame(_BasePandasLikeFrame, Protocol):
@@ -112,6 +111,13 @@ class _ModinDataFrame(_BasePandasLikeFrame, Protocol):
     class _ModinSeries(_BasePandasLikeSeries, Protocol):
         _pandas_class: type[pd.Series[Any]]
 
+    # NOTE: Using `pyspark.sql.DataFrame` creates false positives in overloads when not installed
+    class _PySparkDataFrame(NativeLazyFrame, Protocol):
+        # Not on sqlframe classes
+        # Insane method name that no other framework would clobber
+        # https://github.com/apache/spark/blob/8530444e25b83971da4314c608aa7d763adeceb3/python/pyspark/sql/dataframe.py#L4875
+        def dropDuplicatesWithinWatermark(self, *arg: Any, **kwargs: Any) -> Any: ...  # noqa: N802
+
     _NativePolars: TypeAlias = "pl.DataFrame | pl.LazyFrame | pl.Series"
     _NativeArrow: TypeAlias = "pa.Table | pa.ChunkedArray[Any]"
     _NativeDuckDB: TypeAlias = "duckdb.DuckDBPyRelation"
@@ -124,8 +130,8 @@ class _ModinSeries(_BasePandasLikeSeries, Protocol):
     )
     _NativePandasLike: TypeAlias = "_NativePandasLikeDataFrame |_NativePandasLikeSeries"
     _NativeSQLFrame: TypeAlias = "SQLFrameDataFrame"
-    _NativePySpark: TypeAlias = "pyspark_sql.DataFrame"
-    _NativePySparkConnect: TypeAlias = "PySparkConnectDataFrame"
+    _NativePySpark: TypeAlias = _PySparkDataFrame
+    _NativePySparkConnect: TypeAlias = _PySparkDataFrame
     _NativeSparkLike: TypeAlias = (
         "_NativeSQLFrame | _NativePySpark | _NativePySparkConnect"
     )
@@ -371,8 +377,10 @@ def is_native_dask(obj: Any) -> TypeIs[_NativeDask]:
 
 is_native_duckdb: _Guard[_NativeDuckDB] = is_duckdb_relation
 is_native_sqlframe: _Guard[_NativeSQLFrame] = is_sqlframe_dataframe
-is_native_pyspark: _Guard[_NativePySpark] = is_pyspark_dataframe
-is_native_pyspark_connect: _Guard[_NativePySparkConnect] = is_pyspark_connect_dataframe
+is_native_pyspark = cast("_Guard[_NativePySpark]", is_pyspark_dataframe)
+is_native_pyspark_connect = cast(
+    "_Guard[_NativePySparkConnect]", is_pyspark_connect_dataframe
+)
 
 
 def is_native_pandas(obj: Any) -> TypeIs[_NativePandas]:

narwhals/dataframe.py

@@ -71,8 +71,34 @@
 
     from narwhals._compliant import CompliantDataFrame, CompliantLazyFrame
     from narwhals._compliant.typing import CompliantExprAny, EagerNamespaceAny
+    from narwhals._namespace import (
+        _CuDFDataFrame,
+        _ModinDataFrame,
+        _NativeDask,
+        _NativeDuckDB,
+        _NativePandasLikeDataFrame,
+        _NativeSQLFrame,
+    )
     from narwhals._translate import IntoArrowTable
-    from narwhals._typing import Dask, DuckDB, EagerAllowed, Ibis, IntoBackend, Polars
+    from narwhals._typing import (
+        Dask,
+        DuckDB,
+        EagerAllowed,
+        Ibis,
+        IntoBackend,
+        Polars,
+        _ArrowImpl,
+        _CudfImpl,
+        _DaskImpl,
+        _DuckDBImpl,
+        _EagerAllowedImpl,
+        _LazyAllowedImpl,
+        _ModinImpl,
+        _PandasImpl,
+        _PandasLikeImpl,
+        _PolarsImpl,
+        _SQLFrameImpl,
+    )
     from narwhals.group_by import GroupBy, LazyGroupBy
     from narwhals.typing import (
         AsofJoinStrategy,
@@ -104,10 +130,76 @@
 MultiIndexSelector: TypeAlias = "_MultiIndexSelector[Series[Any]]"
 
 
+class _ImplDescriptor:
+    def __set_name__(self, owner: type[Any], name: str) -> None:
+        self.__name__: str = name
+
+    @overload
+    def __get__(
+        self, instance: DataFrame[pl.DataFrame] | LazyFrame[pl.LazyFrame], owner: Any
+    ) -> _PolarsImpl: ...
+    @overload
+    def __get__(self, instance: BaseFrame[pd.DataFrame], owner: Any) -> _PandasImpl: ...
+    @overload
+    def __get__(self, instance: BaseFrame[_ModinDataFrame], owner: Any) -> _ModinImpl: ...
+
+    @overload  # oof, looks like these two need their names aligned 😅
+    def __get__(self, instance: BaseFrame[_CuDFDataFrame], owner: Any) -> _CudfImpl: ...
+    @overload
+    def __get__(
+        self, instance: BaseFrame[_NativePandasLikeDataFrame], owner: Any
+    ) -> _PandasLikeImpl: ...
+    @overload
+    def __get__(self, instance: BaseFrame[pa.Table], owner: Any) -> _ArrowImpl: ...
+    @overload
+    def __get__(
+        self, instance: BaseFrame[pl.DataFrame | pd.DataFrame | pa.Table], owner: Any
+    ) -> _PolarsImpl | _PandasImpl | _ArrowImpl: ...
+    @overload
+    def __get__(self, instance: LazyFrame[_NativeDuckDB], owner: Any) -> _DuckDBImpl: ...
+    @overload
+    def __get__(
+        self, instance: LazyFrame[_NativeSQLFrame], owner: Any
+    ) -> _SQLFrameImpl: ...
+    @overload
+    def __get__(self, instance: LazyFrame[_NativeDask], owner: Any) -> _DaskImpl: ...
+    @overload
+    def __get__(self, instance: None, owner: Any) -> Self: ...
+    @overload
+    def __get__(self, instance: DataFrame[Any], owner: Any) -> _EagerAllowedImpl: ...
+    @overload
+    def __get__(self, instance: LazyFrame[Any], owner: Any) -> _LazyAllowedImpl: ...
+    def __get__(self, instance: Any | None, owner: Any) -> Any:
+        if instance is None:  # pragma: no cover
+            return self
+        return instance._compliant_frame._implementation
+
+
 class BaseFrame(Generic[_FrameT]):
     _compliant_frame: Any
     _level: Literal["full", "lazy", "interchange"]
 
+    implementation: _ImplDescriptor = _ImplDescriptor()
+    """Return implementation of native frame.
+
+    This can be useful when you need to use special-casing for features outside of
+    Narwhals' scope - for example, when dealing with pandas' Period Dtype.
+
+    Examples:
+        >>> import narwhals as nw
+        >>> import pandas as pd
+        >>> df_native = pd.DataFrame({"a": [1, 2, 3]})
+        >>> df = nw.from_native(df_native)
+        >>> df.implementation
+        <Implementation.PANDAS: 'pandas'>
+        >>> df.implementation.is_pandas()
+        True
+        >>> df.implementation.is_pandas_like()
+        True
+        >>> df.implementation.is_polars()
+        False
+    """
+
     def __native_namespace__(self) -> ModuleType:
         return self._compliant_frame.__native_namespace__()  # type: ignore[no-any-return]
 
@@ -660,29 +752,6 @@ def from_numpy(
         )
         raise ValueError(msg)
 
-    @property
-    def implementation(self) -> Implementation:
-        """Return implementation of native frame.
-
-        This can be useful when you need to use special-casing for features outside of
-        Narwhals' scope - for example, when dealing with pandas' Period Dtype.
-
-        Examples:
-            >>> import narwhals as nw
-            >>> import pandas as pd
-            >>> df_native = pd.DataFrame({"a": [1, 2, 3]})
-            >>> df = nw.from_native(df_native)
-            >>> df.implementation
-            <Implementation.PANDAS: 'pandas'>
-            >>> df.implementation.is_pandas()
-            True
-            >>> df.implementation.is_pandas_like()
-            True
-            >>> df.implementation.is_polars()
-            False
-        """
-        return self._compliant_frame._implementation
-
     def __len__(self) -> int:
         return self._compliant_frame.__len__()
 
@@ -2295,22 +2364,6 @@ def __init__(self, df: Any, *, level: Literal["full", "lazy", "interchange"]) ->
     def __repr__(self) -> str:  # pragma: no cover
         return generate_repr("Narwhals LazyFrame", self.to_native().__repr__())
 
-    @property
-    def implementation(self) -> Implementation:
-        """Return implementation of native frame.
-
-        This can be useful when you need to use special-casing for features outside of
-        Narwhals' scope - for example, when dealing with pandas' Period Dtype.
-
-        Examples:
-            >>> import narwhals as nw
-            >>> import dask.dataframe as dd
-            >>> lf_native = dd.from_dict({"a": [1, 2]}, npartitions=1)
-            >>> nw.from_native(lf_native).implementation
-            <Implementation.DASK: 'dask'>
-        """
-        return self._compliant_frame._implementation
-
     def __getitem__(self, item: str | slice) -> NoReturn:
         msg = "Slicing is not supported on LazyFrame"
         raise TypeError(msg)

pyproject.toml

@@ -74,7 +74,7 @@ typing = [  # keep some of these pinned and bump periodically so there's fewer s
   "sqlframe",
   "polars==1.32.2",
   "uv",
-  "narwhals[ibis]",
+  "narwhals[ibis,dask,modin]",
 ]
 docs = [
   "black",  # required by mkdocstrings_handlers

tests/expr_and_series/is_close_test.py

@@ -11,7 +11,6 @@
 import pytest
 
 import narwhals as nw
-from narwhals._utils import is_eager_allowed
 from narwhals.exceptions import ComputeError, InvalidOperationError
 from tests.conftest import (
     dask_lazy_p1_constructor,
@@ -114,11 +113,8 @@ def test_is_close_series_with_series(
 ) -> None:
     df = nw.from_native(constructor_eager(data), eager_only=True)
     x, y = df["x"], df["y"]
-    backend = df.implementation
-    assert is_eager_allowed(backend)
-
     nulls = nw.new_series(
-        name="nulls", values=[None] * len(x), dtype=nw.Float64(), backend=backend
+        "nulls", [None] * len(x), nw.Float64(), backend=df.implementation
     )
     x = x.zip_with(x != NAN_PLACEHOLDER, x**0.5).zip_with(x != NULL_PLACEHOLDER, nulls)
     y = y.zip_with(y != NAN_PLACEHOLDER, y**0.5).zip_with(y != NULL_PLACEHOLDER, nulls)
@@ -141,11 +137,8 @@ def test_is_close_series_with_scalar(
 ) -> None:
     df = nw.from_native(constructor_eager(data), eager_only=True)
     y = df["y"]
-    backend = df.implementation
-    assert is_eager_allowed(backend)
-
     nulls = nw.new_series(
-        name="nulls", values=[None] * len(y), dtype=nw.Float64(), backend=backend
+        "nulls", [None] * len(y), nw.Float64(), backend=df.implementation
     )
     y = y.zip_with(y != NAN_PLACEHOLDER, y**0.5).zip_with(y != NULL_PLACEHOLDER, nulls)
     result = y.is_close(other, abs_tol=abs_tol, rel_tol=rel_tol, nans_equal=nans_equal)

tests/implementation_test.py

@@ -1,9 +1,27 @@
 from __future__ import annotations
 
+# Using pyright's assert type instead
+# mypy: disable-error-code="assert-type"
+from typing import TYPE_CHECKING, Any, cast
+
 import pytest
 
 import narwhals as nw
 
+if TYPE_CHECKING:
+    from narwhals._typing import (
+        _ArrowImpl,
+        _DaskImpl,
+        _DuckDBImpl,
+        _EagerAllowedImpl,
+        _LazyAllowedImpl,
+        _ModinImpl,
+        _PandasImpl,
+        _PolarsImpl,
+        _SQLFrameImpl,
+    )
+    from narwhals.typing import IntoDataFrame
+
 
 def test_implementation_pandas() -> None:
     pytest.importorskip("pandas")
@@ -52,3 +70,82 @@ def test_implementation_polars() -> None:
 )
 def test_implementation_new(member: str, value: str) -> None:
     assert nw.Implementation(value) is getattr(nw.Implementation, member)
+
+
+if TYPE_CHECKING:
+
+    def test_implementation_typing() -> None:  # noqa: PLR0914
+        import dask.dataframe as dd
+        import modin.pandas as mpd
+        import pandas as pd
+        import polars as pl
+        import pyarrow as pa
+        from typing_extensions import assert_type
+
+        from tests.conftest import (
+            duckdb_lazy_constructor,
+            sqlframe_pyspark_lazy_constructor,
+        )
+
+        data: dict[str, Any] = {"a": [1, 2, 3]}
+        polars_df = nw.from_native(pl.DataFrame(data))
+        polars_ldf = nw.from_native(pl.LazyFrame(data))
+        pandas_df = nw.from_native(pd.DataFrame(data))
+        arrow_df = nw.from_native(pa.table(data))
+        duckdb_ldf = nw.from_native(duckdb_lazy_constructor(data))
+        sqlframe_ldf = nw.from_native(sqlframe_pyspark_lazy_constructor(data))
+
+        polars_impl = polars_df.implementation
+        lazy_polars_impl = polars_ldf.implementation
+        pandas_impl = pandas_df.implementation
+        arrow_impl = arrow_df.implementation
+        duckdb_impl = duckdb_ldf.implementation
+        sqlframe_impl = sqlframe_ldf.implementation
+
+        assert_type(polars_impl, _PolarsImpl)
+        assert_type(lazy_polars_impl, _PolarsImpl)
+        # NOTE: Testing the lazy versions of pandas/pyarrow would require adding overloads to `DataFrame.lazy`
+        # Currently, everything becomes `LazyFrame[Any]`
+        assert_type(pandas_impl, _PandasImpl)
+        assert_type(arrow_impl, _ArrowImpl)
+        assert_type(duckdb_impl, _DuckDBImpl)
+        assert_type(sqlframe_impl, _SQLFrameImpl)
+
+        modin_native = mpd.DataFrame.from_dict(data)
+        modin_df = nw.from_native(modin_native)
+        modin_impl = modin_df.implementation
+        # TODO @dangotbanned: Is this even possible?
+        # - `mypy` won't ever work, treats as `Any`
+        # - `pyright` can resolve `modin_df: narwhals.dataframe.DataFrame[modin.pandas.dataframe.DataFrame]`
+        #   - But we run into variance issues if trying to widen the concrete type again
+        assert_type(modin_impl, _ModinImpl)  # pyright: ignore[reportAssertTypeFailure]
+        # If ^^^ can be fixed, the next one should be removed
+        assert_type(modin_impl, _EagerAllowedImpl)
+
+        # NOTE: Constructor returns `Unknown`
+        dask_native = cast("dd.DataFrame", dd.DataFrame.from_dict(data))
+        dask_ldf = nw.from_native(dask_native)
+        dask_impl = dask_ldf.implementation
+        # NOTE: Same issue as modin
+        assert_type(dask_impl, _DaskImpl)  # pyright: ignore[reportAssertTypeFailure]
+        # If ^^^ can be fixed, the next one should be removed
+        assert_type(dask_impl, _LazyAllowedImpl)
+
+        can_lazyframe_collect_dfs: list[
+            nw.DataFrame[pl.DataFrame]
+            | nw.DataFrame[pd.DataFrame]
+            | nw.DataFrame[pa.Table]
+        ] = [polars_df, pandas_df, arrow_df]
+        can_lazyframe_collect_impl = can_lazyframe_collect_dfs[0].implementation
+        assert_type(can_lazyframe_collect_impl, _PolarsImpl | _PandasImpl | _ArrowImpl)
+
+        very_lost_df = nw.DataFrame.__new__(nw.DataFrame)
+        very_lost_impl = very_lost_df.implementation
+        # TODO @dangotbanned: Is this so bad?
+        # - Currently `DataFrame[Any]` matches the first overload (`_PolarsImpl`)
+        # - That is accepted **everywhere** that uses `IntoBackend`
+        assert_type(very_lost_impl, _EagerAllowedImpl)  # pyright: ignore[reportAssertTypeFailure]
+
+        not_so_lost_df = nw.DataFrame.__new__(nw.DataFrame[IntoDataFrame])
+        not_so_lost_impl = not_so_lost_df.implementation
+        assert_type(not_so_lost_impl, _EagerAllowedImpl)

tests/namespace_test.py

@@ -81,7 +81,7 @@ def test_namespace_from_native_object(constructor: Constructor) -> None:
 def test_namespace_from_native_object_invalid() -> None:
     data = {"a": [1, 2, 3], "b": [4, 5, 6]}
     with pytest.raises(TypeError, match=r"dict"):
-        Namespace.from_native_object(data)  # pyright: ignore[reportCallIssue, reportArgumentType]
+        Namespace.from_native_object(data)  # type: ignore[call-overload]
 
 
 @eager_allowed