chore(typing): Spec out CompliantDataFrame

- Variance of `CompliantSeriesT_co` changed due to how it is used
- Made a lot of the types more generic
  - They can be refined by the implementations
  - Some will handle coverting to narrower types, where that backend demands it
- **Main missing part** is `CompliantExpr`
  - Need to spec `CompliantSeries` first
  - Otherwise, introduces lots of errors for `Polars*`

Author: dangotbanned

narwhals/_compliant/__init__.py

@@ -19,7 +19,7 @@
 from narwhals._compliant.typing import CompliantExprT
 from narwhals._compliant.typing import CompliantFrameT
 from narwhals._compliant.typing import CompliantSeriesOrNativeExprT_co
-from narwhals._compliant.typing import CompliantSeriesT_co
+from narwhals._compliant.typing import CompliantSeriesT
 from narwhals._compliant.typing import EagerDataFrameT
 from narwhals._compliant.typing import EagerSeriesT
 from narwhals._compliant.typing import IntoCompliantExpr
@@ -35,7 +35,7 @@
     "CompliantSelectorNamespace",
     "CompliantSeries",
     "CompliantSeriesOrNativeExprT_co",
-    "CompliantSeriesT_co",
+    "CompliantSeriesT",
     "EagerDataFrame",
     "EagerDataFrameT",
     "EagerExpr",

narwhals/_compliant/dataframe.py

@@ -4,43 +4,152 @@
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterator
+from typing import Literal
 from typing import Mapping
 from typing import Protocol
 from typing import Sequence
+from typing import Sized
 from typing import TypeVar
+from typing import overload
 
-from narwhals._compliant.typing import CompliantSeriesT_co
+from narwhals._compliant.typing import CompliantSeriesT
 from narwhals._compliant.typing import EagerSeriesT
 from narwhals._expression_parsing import evaluate_output_names_and_aliases
 
 if TYPE_CHECKING:
+    from io import BytesIO
+    from pathlib import Path
+
+    import pandas as pd
+    import polars as pl
+    import pyarrow as pa
     from typing_extensions import Self
     from typing_extensions import TypeIs
 
     from narwhals._compliant.expr import EagerExpr
     from narwhals.dtypes import DType
+    from narwhals.typing import SizeUnit
+    from narwhals.typing import _2DArray
+    from narwhals.utils import Implementation
 
 __all__ = ["CompliantDataFrame", "CompliantLazyFrame", "EagerDataFrame"]
 
 T = TypeVar("T")
 
 
-class CompliantDataFrame(Protocol[CompliantSeriesT_co]):
+class CompliantDataFrame(Sized, Protocol[CompliantSeriesT]):
     def __narwhals_dataframe__(self) -> Self: ...
     def __narwhals_namespace__(self) -> Any: ...
-    def simple_select(
-        self, *column_names: str
-    ) -> Self: ...  # `select` where all args are column names.
+    def __array__(self, dtype: Any, *, copy: bool | None) -> _2DArray: ...
+    def __getitem__(self, item: Any) -> CompliantSeriesT | Self: ...
+    def simple_select(self, *column_names: str) -> Self:
+        """`select` where all args are column names."""
+        ...
+
     def aggregate(self, *exprs: Any) -> Self:  # pragma: no cover
-        ...  # `select` where all args are aggregations or literals
-        # (so, no broadcasting is necessary).
+        """`select` where all args are aggregations or literals.
+
+        (so, no broadcasting is necessary).
+        """
+        ...
 
     @property
     def columns(self) -> Sequence[str]: ...
     @property
     def schema(self) -> Mapping[str, DType]: ...
-    def get_column(self, name: str) -> CompliantSeriesT_co: ...
-    def iter_columns(self) -> Iterator[CompliantSeriesT_co]: ...
+    @property
+    def shape(self) -> tuple[int, int]: ...
+    def clone(self) -> Self: ...
+    def collect(
+        self, backend: Implementation | None, **kwargs: Any
+    ) -> CompliantDataFrame[Any]: ...
+    def collect_schema(self) -> Mapping[str, DType]: ...
+    def drop(self, columns: Sequence[str], strict: bool) -> Self: ...  # noqa: FBT001
+    def drop_nulls(self, subset: Sequence[str] | None) -> Self: ...
+    def estimated_size(self, unit: SizeUnit) -> int | float: ...
+    def filter(self, predicate: Any) -> Self: ...
+    def gather_every(self, n: int, offset: int) -> Self: ...
+    def get_column(self, name: str) -> CompliantSeriesT: ...
+    def group_by(self, *keys: str, drop_null_keys: bool) -> Any: ...
+    def head(self, n: int) -> Self: ...
+    def item(self, row: int | None, column: int | str | None) -> Any: ...
+    def iter_columns(self) -> Iterator[CompliantSeriesT]: ...
+    def iter_rows(
+        self, *, named: bool, buffer_size: int
+    ) -> Iterator[tuple[Any, ...]] | Iterator[Mapping[str, Any]]: ...
+    def is_unique(self) -> CompliantSeriesT: ...
+    def join(
+        self: Self,
+        other: Self,
+        *,
+        how: Literal["left", "inner", "cross", "anti", "semi"],
+        left_on: Sequence[str] | None,
+        right_on: Sequence[str] | None,
+        suffix: str,
+    ) -> Self: ...
+    def join_asof(
+        self: Self,
+        other: Self,
+        *,
+        left_on: str | None,
+        right_on: str | None,
+        by_left: Sequence[str] | None,
+        by_right: Sequence[str] | None,
+        strategy: Literal["backward", "forward", "nearest"],
+        suffix: str,
+    ) -> Self: ...
+    def lazy(self, *, backend: Implementation | None) -> CompliantLazyFrame: ...
+    def rename(self, mapping: Mapping[str, str]) -> Self: ...
+    def row(self, index: int) -> tuple[Any, ...]: ...
+    def rows(
+        self, *, named: bool
+    ) -> Sequence[tuple[Any, ...]] | Sequence[Mapping[str, Any]]: ...
+    def sample(
+        self,
+        n: int | None,
+        *,
+        fraction: float | None,
+        with_replacement: bool,
+        seed: int | None,
+    ) -> Self: ...
+    def select(self, *exprs: Any) -> Self: ...
+    def sort(
+        self, *by: str, descending: bool | Sequence[bool], nulls_last: bool
+    ) -> Self: ...
+    def tail(self, n: int) -> Self: ...
+    def to_arrow(self) -> pa.Table: ...
+    def to_numpy(self) -> _2DArray: ...
+    def to_pandas(self) -> pd.DataFrame: ...
+    def to_polars(self) -> pl.DataFrame: ...
+    @overload
+    def to_dict(self, *, as_series: Literal[True]) -> dict[str, CompliantSeriesT]: ...
+    @overload
+    def to_dict(self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ...
+    def to_dict(
+        self, *, as_series: bool
+    ) -> dict[str, CompliantSeriesT] | dict[str, list[Any]]: ...
+    def unique(
+        self,
+        subset: Sequence[str] | None,
+        *,
+        keep: Literal["any", "first", "last", "none"],
+        maintain_order: bool | None,
+    ) -> Self: ...
+    def unpivot(
+        self,
+        on: Sequence[str] | None,
+        index: Sequence[str] | None,
+        variable_name: str,
+        value_name: str,
+    ) -> Self: ...
+    def with_columns(self, *exprs: Any) -> Self: ...
+    def with_row_index(self, name: str) -> Self: ...
+    @overload
+    def write_csv(self, file: None) -> str: ...
+    @overload
+    def write_csv(self, file: str | Path | BytesIO) -> None: ...
+    def write_csv(self, file: str | Path | BytesIO | None) -> str | None: ...
+    def write_parquet(self, file: str | Path | BytesIO) -> None: ...
 
 
 class CompliantLazyFrame(Protocol):

narwhals/_compliant/typing.py

@@ -25,13 +25,11 @@
     "CompliantDataFrameT",
     "CompliantFrameT",
     "CompliantLazyFrameT",
-    "CompliantSeriesT_co",
+    "CompliantSeriesT",
     "IntoCompliantExpr",
 ]
 NativeExprT_co = TypeVar("NativeExprT_co", bound="NativeExpr", covariant=True)
-CompliantSeriesT_co = TypeVar(
-    "CompliantSeriesT_co", bound="CompliantSeries", covariant=True
-)
+CompliantSeriesT = TypeVar("CompliantSeriesT", bound="CompliantSeries")
 CompliantSeriesOrNativeExprT_co = TypeVar(
     "CompliantSeriesOrNativeExprT_co",
     bound="CompliantSeries | NativeExpr",