Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 71 additions & 83 deletions narwhals/_arrow/dataframe.py

Large diffs are not rendered by default.

2 changes: 2 additions & 0 deletions narwhals/_compliant/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from narwhals._compliant.typing import EagerDataFrameT
from narwhals._compliant.typing import EagerSeriesT
from narwhals._compliant.typing import IntoCompliantExpr
from narwhals._compliant.typing import NativeFrameT_co

__all__ = [
"CompliantDataFrame",
Expand All @@ -48,4 +49,5 @@
"IntoCompliantExpr",
"LazyExpr",
"LazySelectorNamespace",
"NativeFrameT_co",
]
18 changes: 14 additions & 4 deletions narwhals/_compliant/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,13 @@
T = TypeVar("T")


class CompliantDataFrame(Sized, Protocol[CompliantSeriesT, CompliantExprT_contra]):
class CompliantDataFrame(
_StoresNative[NativeFrameT_co],
Sized,
Protocol[CompliantSeriesT, CompliantExprT_contra, NativeFrameT_co],
):
_native_frame: Any

def __narwhals_dataframe__(self) -> Self: ...
def __narwhals_namespace__(self) -> Any: ...
def __array__(self, dtype: Any, *, copy: bool | None) -> _2DArray: ...
Expand All @@ -60,6 +66,10 @@ def aggregate(self, *exprs: CompliantExprT_contra) -> Self:
"""
return self.select(*exprs)

@property
def native(self) -> NativeFrameT_co:
return self._native_frame # type: ignore[no-any-return]

@property
def columns(self) -> Sequence[str]: ...
@property
Expand All @@ -69,7 +79,7 @@ def shape(self) -> tuple[int, int]: ...
def clone(self) -> Self: ...
def collect(
self, backend: Implementation | None, **kwargs: Any
) -> CompliantDataFrame[Any, Any]: ...
) -> CompliantDataFrame[Any, Any, Any]: ...
def collect_schema(self) -> Mapping[str, DType]: ...
def drop(self, columns: Sequence[str], *, strict: bool) -> Self: ...
def drop_nulls(self, subset: Sequence[str] | None) -> Self: ...
Expand Down Expand Up @@ -195,7 +205,7 @@ def schema(self) -> Mapping[str, DType]: ...
def _iter_columns(self) -> Iterator[Any]: ...
def collect(
self, backend: Implementation | None, **kwargs: Any
) -> CompliantDataFrame[Any, Any]: ...
) -> CompliantDataFrame[Any, Any, Any]: ...
def collect_schema(self) -> Mapping[str, DType]: ...
def drop(self, columns: Sequence[str], *, strict: bool) -> Self: ...
def drop_nulls(self, subset: Sequence[str] | None) -> Self: ...
Expand Down Expand Up @@ -252,7 +262,7 @@ def with_row_index(self, name: str) -> Self: ...


class EagerDataFrame(
CompliantDataFrame[EagerSeriesT, EagerExprT_contra],
CompliantDataFrame[EagerSeriesT, EagerExprT_contra, NativeFrameT_co],
CompliantLazyFrame[EagerExprT_contra, NativeFrameT_co],
Protocol[EagerSeriesT, EagerExprT_contra, NativeFrameT_co],
):
Expand Down
6 changes: 3 additions & 3 deletions narwhals/_compliant/selectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,9 +66,9 @@
SeriesT = TypeVar("SeriesT", bound="CompliantSeries")
ExprT = TypeVar("ExprT", bound="NativeExpr")
FrameT = TypeVar(
"FrameT", bound="CompliantDataFrame[Any, Any] | CompliantLazyFrame[Any, Any]"
"FrameT", bound="CompliantDataFrame[Any, Any, Any] | CompliantLazyFrame[Any, Any]"
)
DataFrameT = TypeVar("DataFrameT", bound="CompliantDataFrame[Any, Any]")
DataFrameT = TypeVar("DataFrameT", bound="CompliantDataFrame[Any, Any, Any]")
LazyFrameT = TypeVar("LazyFrameT", bound="CompliantLazyFrame[Any, Any]")
SelectorOrExpr: TypeAlias = (
"CompliantSelector[FrameT, SeriesOrExprT] | CompliantExpr[FrameT, SeriesOrExprT]"
Expand Down Expand Up @@ -311,7 +311,7 @@ def __repr__(self: Self) -> str: # pragma: no cover


def _eval_lhs_rhs(
df: CompliantDataFrame[Any, Any] | CompliantLazyFrame[Any, Any],
df: CompliantDataFrame[Any, Any, Any] | CompliantLazyFrame[Any, Any],
lhs: CompliantExpr[Any, Any],
rhs: CompliantExpr[Any, Any],
) -> tuple[Sequence[str], Sequence[str]]:
Expand Down
9 changes: 6 additions & 3 deletions narwhals/_compliant/typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
"CompliantLazyFrameT",
"CompliantSeriesT",
"IntoCompliantExpr",
"NativeFrameT_co",
]
NativeExprT_co = TypeVar("NativeExprT_co", bound="NativeExpr", covariant=True)
CompliantSeriesT = TypeVar("CompliantSeriesT", bound="CompliantSeries")
Expand All @@ -36,12 +37,14 @@
bound="CompliantSeries | NativeExpr",
covariant=True,
)

NativeFrameT_co = TypeVar("NativeFrameT_co", bound="NativeFrame", covariant=True)
CompliantFrameT = TypeVar(
"CompliantFrameT", bound="CompliantDataFrame[Any, Any] | CompliantLazyFrame[Any, Any]"
"CompliantFrameT",
bound="CompliantDataFrame[Any, Any, Any] | CompliantLazyFrame[Any, Any]",
)
CompliantDataFrameT = TypeVar(
"CompliantDataFrameT", bound="CompliantDataFrame[Any, Any, Any]"
)
CompliantDataFrameT = TypeVar("CompliantDataFrameT", bound="CompliantDataFrame[Any, Any]")
CompliantLazyFrameT = TypeVar("CompliantLazyFrameT", bound="CompliantLazyFrame[Any, Any]")
IntoCompliantExpr: TypeAlias = "CompliantExpr[CompliantFrameT, CompliantSeriesOrNativeExprT_co] | CompliantSeriesOrNativeExprT_co"
CompliantExprT = TypeVar("CompliantExprT", bound="CompliantExpr[Any, Any]")
Expand Down
2 changes: 1 addition & 1 deletion narwhals/_dask/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ def collect(
self: Self,
backend: Implementation | None,
**kwargs: Any,
) -> CompliantDataFrame[Any, Any]:
) -> CompliantDataFrame[Any, Any, Any]:
result = self._native_frame.compute(**kwargs)

if backend is None or backend is Implementation.PANDAS:
Expand Down
2 changes: 1 addition & 1 deletion narwhals/_duckdb/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def collect(
self: Self,
backend: ModuleType | Implementation | str | None,
**kwargs: Any,
) -> CompliantDataFrame[Any, Any]:
) -> CompliantDataFrame[Any, Any, Any]:
if backend is None or backend is Implementation.PYARROW:
import pyarrow as pa # ignore-banned-import

Expand Down
2 changes: 1 addition & 1 deletion narwhals/_expression_parsing.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ def extract_compliant(

def evaluate_output_names_and_aliases(
expr: CompliantExpr[Any, Any],
df: CompliantDataFrame[Any, Any] | CompliantLazyFrame[Any, Any],
df: CompliantDataFrame[Any, Any, Any] | CompliantLazyFrame[Any, Any],
exclude: Sequence[str],
) -> tuple[Sequence[str], Sequence[str]]:
output_names = expr._evaluate_output_names(df)
Expand Down
Loading
Loading