feat(typing): Make Implementation less opaque (#3016)

Author: dangotbanned

.github/workflows/typing.yml

@@ -31,7 +31,7 @@ jobs:
         run: uv venv .venv
       - name: install-reqs
         # TODO: add more dependencies/backends incrementally
-        run: uv pip install -e ".[pyspark]" --group core --group typing
+        run: uv pip install -e ".[pyspark]" --group core --group typing-ci
       - name: show-deps
         run: uv pip freeze
       - name: Run mypy and pyright

Makefile

@@ -20,6 +20,6 @@ help:  ## Display this help screen
 
 .PHONY: typing
 typing: ## Run typing checks
-	$(VENV_BIN)/uv pip install -e . --group typing
+	$(VENV_BIN)/uv pip install -e . --group typing-ci
 	$(VENV_BIN)/pyright
 	$(VENV_BIN)/mypy

docs/api-reference/dataframe.md

@@ -58,3 +58,4 @@
         - write_parquet
       show_source: false
       show_bases: false
+      inherited_members: true

docs/api-reference/lazyframe.md

@@ -34,3 +34,4 @@
       show_root_heading: false
       show_source: false
       show_bases: false
+      inherited_members: true

narwhals/_namespace.py

@@ -7,9 +7,9 @@
     Any,
     Callable,
     Generic,
-    Literal,
     Protocol,
     TypeVar,
+    cast,
     overload,
 )
 
@@ -37,8 +37,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
@@ -68,30 +66,33 @@
     _Guard: TypeAlias = "Callable[[Any], TypeIs[T]]"
 
     EagerAllowedNamespace: TypeAlias = "Namespace[PandasLikeNamespace] | Namespace[ArrowNamespace] | Namespace[PolarsNamespace]"
+    Incomplete: TypeAlias = Any
 
     class _BasePandasLike(Sized, Protocol):
         index: Any
         """`mypy` doesn't like the asymmetric `property` setter in `pandas`."""
 
         def __getitem__(self, key: Any, /) -> Any: ...
-        def __mul__(self, other: float | Collection[float] | Self) -> Self: ...
-        def __floordiv__(self, other: float | Collection[float] | Self) -> Self: ...
+        def __mul__(self, other: float | Collection[float] | Self, /) -> Self: ...
+        def __floordiv__(self, other: float | Collection[float] | Self, /) -> Self: ...
         @property
         def loc(self) -> Any: ...
         @property
         def shape(self) -> tuple[int, ...]: ...
         def set_axis(self, labels: Any, *, axis: Any = ..., copy: bool = ...) -> Self: ...
         def copy(self, deep: bool = ...) -> Self: ...  # noqa: FBT001
-        def rename(self, *args: Any, inplace: Literal[False], **kwds: Any) -> Self:
-            """`inplace=False` is required to avoid (incorrect?) default overloads."""
-            ...
+        def rename(self, *args: Any, **kwds: Any) -> Self | Incomplete:
+            """`mypy` & `pyright` disagree on overloads.
+
+            `Incomplete` used to fix [more important issue](https://github.com/narwhals-dev/narwhals/pull/3016#discussion_r2296139744).
+            """
 
     class _BasePandasLikeFrame(NativeDataFrame, _BasePandasLike, Protocol): ...
 
     class _BasePandasLikeSeries(NativeSeries, _BasePandasLike, Protocol):
-        def where(self, cond: Any, other: Any = ..., **kwds: Any) -> Any: ...
+        def where(self, cond: Any, other: Any = ..., /) -> Self | Incomplete: ...
 
-    class _NativeDask(Protocol):
+    class _NativeDask(NativeLazyFrame, Protocol):
         _partition_type: type[pd.DataFrame]
 
     class _CuDFDataFrame(_BasePandasLikeFrame, Protocol):
@@ -112,6 +113,12 @@ 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):
+        # Arbitrary method that `sqlframe` doesn't have and unlikely to appear anywhere else
+        # 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 +131,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 +378,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/_pandas_like/utils.py

@@ -3,7 +3,7 @@
 import functools
 import operator
 import re
-from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar, cast
 
 import pandas as pd
 
@@ -202,8 +202,10 @@ def rename(
     if implementation is Implementation.PANDAS and (
         implementation._backend_version() >= (3,)
     ):  # pragma: no cover
-        return obj.rename(*args, **kwargs, inplace=False)
-    return obj.rename(*args, **kwargs, copy=False, inplace=False)
+        result = obj.rename(*args, **kwargs, inplace=False)
+    else:
+        result = obj.rename(*args, **kwargs, copy=False, inplace=False)
+    return cast("NativeNDFrameT", result)  # type: ignore[redundant-cast]
 
 
 @functools.lru_cache(maxsize=16)

narwhals/_utils.py

@@ -73,14 +73,40 @@
         NativeSeriesT_co,
     )
     from narwhals._compliant.typing import EvalNames, NativeDataFrameT, NativeLazyFrameT
-    from narwhals._namespace import Namespace
+    from narwhals._namespace import (
+        Namespace,
+        _NativeArrow,
+        _NativeCuDF,
+        _NativeDask,
+        _NativeDuckDB,
+        _NativeIbis,
+        _NativeModin,
+        _NativePandas,
+        _NativePandasLike,
+        _NativePolars,
+        _NativePySpark,
+        _NativePySparkConnect,
+        _NativeSQLFrame,
+    )
     from narwhals._translate import ArrowStreamExportable, IntoArrowTable, ToNarwhalsT_co
     from narwhals._typing import (
         Backend,
         IntoBackend,
+        _ArrowImpl,
+        _CudfImpl,
+        _DaskImpl,
+        _DuckDBImpl,
         _EagerAllowedImpl,
+        _IbisImpl,
         _LazyAllowedImpl,
         _LazyFrameCollectImpl,
+        _ModinImpl,
+        _PandasImpl,
+        _PandasLikeImpl,
+        _PolarsImpl,
+        _PySparkConnectImpl,
+        _PySparkImpl,
+        _SQLFrameImpl,
     )
     from narwhals.dataframe import DataFrame, LazyFrame
     from narwhals.dtypes import DType
@@ -141,7 +167,7 @@ def columns(self) -> Sequence[str]: ...
 _Constructor: TypeAlias = "Callable[Concatenate[_T, P], R2]"
 
 
-class _StoresNative(Protocol[NativeT_co]):  # noqa: PYI046
+class _StoresNative(Protocol[NativeT_co]):
     """Provides access to a native object.
 
     Native objects have types like:
@@ -2034,3 +2060,91 @@ def deep_attrgetter(attr: str, *nested: str) -> attrgetter[Any]:
 def deep_getattr(obj: Any, name_1: str, *nested: str) -> Any:
     """Perform a nested attribute lookup on `obj`."""
     return deep_attrgetter(name_1, *nested)(obj)
+
+
+class Compliant(
+    _StoresNative[NativeT_co], _StoresImplementation, Protocol[NativeT_co]
+): ...
+
+
+class Narwhals(Protocol[NativeT_co]):
+    """Minimal *Narwhals-level* protocol.
+
+    Provides access to a compliant object:
+
+        obj: Narwhals[NativeT_co]]
+        compliant: Compliant[NativeT_co] = obj._compliant
+
+    Which itself exposes:
+
+        implementation: Implementation = compliant.implementation
+        native: NativeT_co = compliant.native
+
+    This interface is used for revealing which `Implementation` member is associated with **either**:
+    - One or more [nominal] native type(s)
+    - One or more [structural] type(s)
+      - where the true native type(s) are [assignable to] *at least* one of them
+
+    These relationships are defined in the `@overload`s of `_Implementation.__get__(...)`.
+
+    [nominal]: https://typing.python.org/en/latest/spec/glossary.html#term-nominal
+    [structural]: https://typing.python.org/en/latest/spec/glossary.html#term-structural
+    [assignable to]: https://typing.python.org/en/latest/spec/glossary.html#term-assignable
+    """
+
+    @property
+    def _compliant(self) -> Compliant[NativeT_co]: ...
+
+
+class _Implementation:
+    """Descriptor for matching an opaque `Implementation` on a generic class.
+
+    Based on [pyright comment](https://github.com/microsoft/pyright/issues/3071#issuecomment-1043978070)
+    """
+
+    def __set_name__(self, owner: type[Any], name: str) -> None:
+        self.__name__: str = name
+
+    @overload
+    def __get__(self, instance: Narwhals[_NativePolars], owner: Any) -> _PolarsImpl: ...
+    @overload
+    def __get__(self, instance: Narwhals[_NativePandas], owner: Any) -> _PandasImpl: ...
+    @overload
+    def __get__(self, instance: Narwhals[_NativeModin], owner: Any) -> _ModinImpl: ...
+    @overload  # TODO @dangotbanned: Rename `_typing` `*Cudf*` aliases to `*CuDF*`
+    def __get__(self, instance: Narwhals[_NativeCuDF], owner: Any) -> _CudfImpl: ...
+    @overload
+    def __get__(
+        self, instance: Narwhals[_NativePandasLike], owner: Any
+    ) -> _PandasLikeImpl: ...
+    @overload
+    def __get__(self, instance: Narwhals[_NativeArrow], owner: Any) -> _ArrowImpl: ...
+    @overload
+    def __get__(
+        self, instance: Narwhals[_NativePolars | _NativeArrow | _NativePandas], owner: Any
+    ) -> _PolarsImpl | _PandasImpl | _ArrowImpl: ...
+    @overload
+    def __get__(self, instance: Narwhals[_NativeDuckDB], owner: Any) -> _DuckDBImpl: ...
+    @overload
+    def __get__(
+        self, instance: Narwhals[_NativeSQLFrame], owner: Any
+    ) -> _SQLFrameImpl: ...
+    @overload
+    def __get__(self, instance: Narwhals[_NativeDask], owner: Any) -> _DaskImpl: ...
+    @overload
+    def __get__(self, instance: Narwhals[_NativeIbis], owner: Any) -> _IbisImpl: ...
+    @overload
+    def __get__(
+        self, instance: Narwhals[_NativePySpark | _NativePySparkConnect], owner: Any
+    ) -> _PySparkImpl | _PySparkConnectImpl: ...
+    # NOTE: https://docs.python.org/3/howto/descriptor.html#invocation-from-a-class
+    @overload
+    def __get__(self, instance: None, owner: type[Narwhals[Any]]) -> Self: ...
+    @overload
+    def __get__(
+        self, instance: DataFrame[Any] | Series[Any], owner: Any
+    ) -> _EagerAllowedImpl: ...
+    @overload
+    def __get__(self, instance: LazyFrame[Any], owner: Any) -> _LazyAllowedImpl: ...
+    def __get__(self, instance: Narwhals[Any] | None, owner: Any) -> Any:
+        return self if instance is None else instance._compliant._implementation