diff --git a/narwhals/_compliant/expr.py b/narwhals/_compliant/expr.py index 910b110a55..4f7169b632 100644 --- a/narwhals/_compliant/expr.py +++ b/narwhals/_compliant/expr.py @@ -168,7 +168,9 @@ def replace_strict( *, return_dtype: DType | type[DType] | None, ) -> Self: ... - def over(self: Self, keys: Sequence[str], order_by: Sequence[str] | None) -> Self: ... + def over( + self, partition_by: Sequence[str], order_by: Sequence[str] | None + ) -> Self: ... def sample( self, n: int | None, diff --git a/narwhals/_polars/expr.py b/narwhals/_polars/expr.py index b498b94f6c..de6c4cdd2c 100644 --- a/narwhals/_polars/expr.py +++ b/narwhals/_polars/expr.py @@ -4,6 +4,7 @@ from typing import Any from typing import Callable from typing import Literal +from typing import Mapping from typing import Sequence import polars as pl @@ -18,6 +19,8 @@ from narwhals._expression_parsing import ExprKind from narwhals._expression_parsing import ExprMetadata + from narwhals._polars.dataframe import Method + from narwhals._polars.namespace import PolarsNamespace from narwhals.dtypes import DType from narwhals.utils import Version @@ -61,7 +64,7 @@ def _renamed_min_periods(self, min_samples: int, /) -> dict[str, Any]: name = "min_periods" if self._backend_version < (1, 21, 0) else "min_samples" return {name: min_samples} - def cast(self: Self, dtype: DType) -> Self: + def cast(self, dtype: DType | type[DType]) -> Self: dtype_pl = narwhals_to_native_dtype(dtype, self._version, self._backend_version) return self._with_native(self.native.cast(dtype_pl)) @@ -96,9 +99,7 @@ def is_nan(self: Self) -> Self: native = pl.when(self.native.is_not_null()).then(self.native.is_nan()) return self._with_native(native) - def over( - self: Self, partition_by: Sequence[str], order_by: Sequence[str] | None - ) -> Self: + def over(self, partition_by: Sequence[str], order_by: Sequence[str] | None) -> Self: if self._backend_version < (1, 9): if order_by: msg = "`order_by` in Polars requires version 1.10 or greater" @@ -147,7 +148,7 @@ def rolling_mean( return self._with_native(native) def map_batches( - self: Self, function: Callable[..., Self], return_dtype: DType | None + self, function: Callable[[Any], Any], return_dtype: DType | type[DType] | None ) -> Self: return_dtype_pl = ( narwhals_to_native_dtype(return_dtype, self._version, self._backend_version) @@ -158,7 +159,11 @@ def map_batches( return self._with_native(native) def replace_strict( - self: Self, old: Sequence[Any], new: Sequence[Any], *, return_dtype: DType | None + self, + old: Sequence[Any] | Mapping[Any, Any], + new: Sequence[Any], + *, + return_dtype: DType | type[DType] | None, ) -> Self: if self._backend_version < (1,): msg = f"`replace_strict` is only available in Polars>=1.0, found version {self._backend_version}" @@ -226,6 +231,14 @@ def cum_count(self: Self, *, reverse: bool) -> Self: result = self.native.cum_count(reverse=reverse) return self._with_native(result) + def __narwhals_expr__(self) -> None: ... + def __narwhals_namespace__(self) -> PolarsNamespace: # pragma: no cover + from narwhals._polars.namespace import PolarsNamespace + + return PolarsNamespace( + backend_version=self._backend_version, version=self._version + ) + @property def dt(self: Self) -> PolarsExprDateTimeNamespace: return PolarsExprDateTimeNamespace(self) @@ -250,6 +263,59 @@ def list(self: Self) -> PolarsExprListNamespace: def struct(self: Self) -> PolarsExprStructNamespace: return PolarsExprStructNamespace(self) + # CompliantExpr + _alias_output_names: Any + _evaluate_output_names: Any + _is_multi_output_unnamed: Any + __call__: Any + from_column_names: Any + from_column_indices: Any + + # Polars + abs: Method[Self] + all: Method[Self] + any: Method[Self] + alias: Method[Self] + arg_max: Method[Self] + arg_min: Method[Self] + arg_true: Method[Self] + count: Method[Self] + cum_max: Method[Self] + cum_min: Method[Self] + cum_prod: Method[Self] + cum_sum: Method[Self] + diff: Method[Self] + drop_nulls: Method[Self] + fill_null: Method[Self] + gather_every: Method[Self] + head: Method[Self] + is_finite: Method[Self] + is_first_distinct: Method[Self] + is_in: Method[Self] + is_last_distinct: Method[Self] + is_null: Method[Self] + is_unique: Method[Self] + len: Method[Self] + max: Method[Self] + mean: Method[Self] + median: Method[Self] + min: Method[Self] + mode: Method[Self] + n_unique: Method[Self] + null_count: Method[Self] + quantile: Method[Self] + rank: Method[Self] + round: Method[Self] + sample: Method[Self] + shift: Method[Self] + skew: Method[Self] + std: Method[Self] + sum: Method[Self] + sort: Method[Self] + tail: Method[Self] + unique: Method[Self] + var: Method[Self] + class PolarsExprDateTimeNamespace: def __init__(self: Self, expr: PolarsExpr) -> None: diff --git a/narwhals/_polars/namespace.py b/narwhals/_polars/namespace.py index 3ea8f945f2..eb58e41545 100644 --- a/narwhals/_polars/namespace.py +++ b/narwhals/_polars/namespace.py @@ -24,7 +24,6 @@ from datetime import timezone from typing_extensions import Self - from typing_extensions import TypeAlias from narwhals._compliant import CompliantSelectorNamespace from narwhals._compliant import CompliantWhen @@ -39,8 +38,6 @@ from narwhals.utils import Version from narwhals.utils import _FullContext - Incomplete: TypeAlias = Any - class PolarsNamespace: all: Method[PolarsExpr] @@ -51,8 +48,10 @@ class PolarsNamespace: sum_horizontal: Method[PolarsExpr] min_horizontal: Method[PolarsExpr] max_horizontal: Method[PolarsExpr] - # NOTE: `PolarsSeries`, `PolarsExpr` still have gaps - when: Method[CompliantWhen[PolarsDataFrame, Incomplete, Incomplete]] + + # NOTE: `pyright` accepts, `mypy` doesn't highlight the issue + # error: Type argument "PolarsExpr" of "CompliantWhen" must be a subtype of "CompliantExpr[Any, Any]" + when: Method[CompliantWhen[PolarsDataFrame, PolarsSeries, PolarsExpr]] # type: ignore[type-var] def __init__( self: Self, *, backend_version: tuple[int, ...], version: Version @@ -231,10 +230,12 @@ def concat_str( # 1. Others have lots of private stuff for code reuse # i. None of that is useful here # 2. We don't have a `PolarsSelector` abstraction, and just use `PolarsExpr` - # 3. `PolarsExpr` still has it's own gaps in the spec @property - def selectors(self: Self) -> CompliantSelectorNamespace[Any, Any]: - return cast("CompliantSelectorNamespace[Any, Any]", PolarsSelectorNamespace(self)) + def selectors(self) -> CompliantSelectorNamespace[PolarsDataFrame, PolarsSeries]: + return cast( + "CompliantSelectorNamespace[PolarsDataFrame, PolarsSeries]", + PolarsSelectorNamespace(self), + ) class PolarsSelectorNamespace: diff --git a/narwhals/_polars/series.py b/narwhals/_polars/series.py index 50481e50e4..162fe8a7fb 100644 --- a/narwhals/_polars/series.py +++ b/narwhals/_polars/series.py @@ -3,6 +3,8 @@ from typing import TYPE_CHECKING from typing import Any from typing import Iterable +from typing import Iterator +from typing import Mapping from typing import Sequence from typing import cast from typing import overload @@ -22,8 +24,11 @@ from types import ModuleType from typing import TypeVar + import pandas as pd from typing_extensions import Self + from narwhals._arrow.typing import ArrowArray + from narwhals._polars.dataframe import Method from narwhals._polars.dataframe import PolarsDataFrame from narwhals._polars.expr import PolarsExpr from narwhals._polars.namespace import PolarsNamespace @@ -177,12 +182,16 @@ def __getitem__( ) -> Any | Self: return self._from_native_object(self.native.__getitem__(item)) - def cast(self: Self, dtype: DType) -> Self: + def cast(self: Self, dtype: DType | type[DType]) -> Self: dtype_pl = narwhals_to_native_dtype(dtype, self._version, self._backend_version) return self._with_native(self.native.cast(dtype_pl)) def replace_strict( - self: Self, old: Sequence[Any], new: Sequence[Any], *, return_dtype: DType | None + self: Self, + old: Sequence[Any] | Mapping[Any, Any], + new: Sequence[Any], + *, + return_dtype: DType | type[DType] | None, ) -> Self: ser = self.native dtype = ( @@ -586,14 +595,74 @@ def str(self: Self) -> PolarsSeriesStringNamespace: def cat(self: Self) -> PolarsSeriesCatNamespace: return PolarsSeriesCatNamespace(self) - @property - def list(self: Self) -> PolarsSeriesListNamespace: - return PolarsSeriesListNamespace(self) - @property def struct(self: Self) -> PolarsSeriesStructNamespace: return PolarsSeriesStructNamespace(self) + __iter__: Method[Iterator[Any]] + __floordiv__: Method[Self] + __mod__: Method[Self] + __rand__: Method[Self] + __rfloordiv__: Method[Self] + __rmod__: Method[Self] + __ror__: Method[Self] + __rtruediv__: Method[Self] + __truediv__: Method[Self] + abs: Method[Self] + all: Method[bool] + any: Method[bool] + arg_max: Method[int] + arg_min: Method[int] + arg_true: Method[Self] + clip: Method[Self] + count: Method[int] + cum_max: Method[Self] + cum_min: Method[Self] + cum_prod: Method[Self] + cum_sum: Method[Self] + diff: Method[Self] + drop_nulls: Method[Self] + fill_null: Method[Self] + filter: Method[Self] + gather_every: Method[Self] + head: Method[Self] + is_between: Method[Self] + is_finite: Method[Self] + is_first_distinct: Method[Self] + is_in: Method[Self] + is_last_distinct: Method[Self] + is_null: Method[Self] + is_sorted: Method[bool] + is_unique: Method[Self] + item: Method[Any] + len: Method[int] + max: Method[Any] + mean: Method[float] + min: Method[Any] + mode: Method[Self] + n_unique: Method[int] + null_count: Method[int] + quantile: Method[float] + rank: Method[Self] + round: Method[Self] + sample: Method[Self] + shift: Method[Self] + skew: Method[float | None] + std: Method[float] + sum: Method[float] + tail: Method[Self] + to_arrow: Method[ArrowArray] + to_frame: Method[PolarsDataFrame] + to_list: Method[list[Any]] + to_pandas: Method[pd.Series[Any]] + unique: Method[Self] + var: Method[float] + zip_with: Method[Self] + + @property + def list(self: Self) -> PolarsSeriesListNamespace: + return PolarsSeriesListNamespace(self) + class PolarsSeriesDateTimeNamespace: def __init__(self: Self, series: PolarsSeries) -> None: