From 859e0d6fb6467f3460b79e5b93012b446f35253e Mon Sep 17 00:00:00 2001 From: Marco Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Date: Sat, 23 Nov 2024 10:53:10 +0000 Subject: [PATCH] make generic again --- narwhals/stable/v1/__init__.py | 76 +++++++++++++++++++--------------- narwhals/translate.py | 12 +++--- 2 files changed, 48 insertions(+), 40 deletions(-) diff --git a/narwhals/stable/v1/__init__.py b/narwhals/stable/v1/__init__.py index 5b28d70b3..9e4432ec4 100644 --- a/narwhals/stable/v1/__init__.py +++ b/narwhals/stable/v1/__init__.py @@ -74,7 +74,6 @@ from narwhals.dtypes import DType from narwhals.functions import ArrowStreamExportable from narwhals.typing import IntoExpr - from narwhals.typing import IntoSeries T = TypeVar("T") @@ -92,7 +91,7 @@ class DataFrame(NwDataFrame[IntoDataFrameT]): # annotations are correct. @property - def _series(self) -> type[Series]: + def _series(self) -> type[Series[Any]]: return Series @property @@ -106,23 +105,23 @@ def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ... @overload def __getitem__(self, item: tuple[slice, Sequence[int]]) -> Self: ... @overload - def __getitem__(self, item: tuple[Sequence[int], str]) -> Series: ... # type: ignore[overload-overlap] + def __getitem__(self, item: tuple[Sequence[int], str]) -> Series[Any]: ... # type: ignore[overload-overlap] @overload - def __getitem__(self, item: tuple[slice, str]) -> Series: ... # type: ignore[overload-overlap] + def __getitem__(self, item: tuple[slice, str]) -> Series[Any]: ... # type: ignore[overload-overlap] @overload def __getitem__(self, item: tuple[Sequence[int], Sequence[str]]) -> Self: ... @overload def __getitem__(self, item: tuple[slice, Sequence[str]]) -> Self: ... @overload - def __getitem__(self, item: tuple[Sequence[int], int]) -> Series: ... # type: ignore[overload-overlap] + def __getitem__(self, item: tuple[Sequence[int], int]) -> Series[Any]: ... # type: ignore[overload-overlap] @overload - def __getitem__(self, item: tuple[slice, int]) -> Series: ... # type: ignore[overload-overlap] + def __getitem__(self, item: tuple[slice, int]) -> Series[Any]: ... # type: ignore[overload-overlap] @overload def __getitem__(self, item: Sequence[int]) -> Self: ... @overload - def __getitem__(self, item: str) -> Series: ... # type: ignore[overload-overlap] + def __getitem__(self, item: str) -> Series[Any]: ... # type: ignore[overload-overlap] @overload def __getitem__(self, item: Sequence[str]) -> Self: ... @@ -188,14 +187,16 @@ def lazy(self) -> LazyFrame[Any]: # Not sure what mypy is complaining about, probably some fancy # thing that I need to understand category theory for @overload # type: ignore[override] - def to_dict(self, *, as_series: Literal[True] = ...) -> dict[str, Series]: ... + def to_dict(self, *, as_series: Literal[True] = ...) -> dict[str, Series[Any]]: ... @overload def to_dict(self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ... @overload - def to_dict(self, *, as_series: bool) -> dict[str, Series] | dict[str, list[Any]]: ... + def to_dict( + self, *, as_series: bool + ) -> dict[str, Series[Any]] | dict[str, list[Any]]: ... def to_dict( self, *, as_series: bool = True - ) -> dict[str, Series] | dict[str, list[Any]]: + ) -> dict[str, Series[Any]] | dict[str, list[Any]]: """Convert DataFrame to a dictionary mapping column name to values. Arguments: @@ -242,7 +243,7 @@ def to_dict( """ return super().to_dict(as_series=as_series) # type: ignore[return-value] - def is_duplicated(self: Self) -> Series: + def is_duplicated(self: Self) -> Series[Any]: r"""Get a mask of all duplicated rows in this DataFrame. Returns: @@ -292,7 +293,7 @@ def is_duplicated(self: Self) -> Series: """ return super().is_duplicated() # type: ignore[return-value] - def is_unique(self: Self) -> Series: + def is_unique(self: Self) -> Series[Any]: r"""Get a mask of all unique rows in this DataFrame. Returns: @@ -410,7 +411,7 @@ def _l1_norm(self: Self) -> Self: return self.select(all()._l1_norm()) -class Series(NwSeries[Any]): +class Series(NwSeries[IntoSeriesT]): """Narwhals Series, backed by a native series. The native series might be pandas.Series, polars.Series, ... @@ -1150,7 +1151,7 @@ def _stableify(obj: NwDataFrame[IntoFrameT]) -> DataFrame[IntoFrameT]: ... @overload def _stableify(obj: NwLazyFrame[IntoFrameT]) -> LazyFrame[IntoFrameT]: ... @overload -def _stableify(obj: NwSeries[IntoSeriesT]) -> Series: ... +def _stableify(obj: NwSeries[IntoSeriesT]) -> Series[IntoSeriesT]: ... @overload def _stableify(obj: NwExpr) -> Expr: ... @overload @@ -1163,7 +1164,7 @@ def _stableify( | NwSeries[IntoSeriesT] | NwExpr | Any, -) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series | Expr | Any: +) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series[Any] | Expr | Any: if isinstance(obj, NwDataFrame): return DataFrame( obj._compliant_frame, @@ -1193,7 +1194,7 @@ def from_native( eager_or_interchange_only: Literal[True], series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[IntoDataFrameT] | Series: ... +) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload @@ -1205,7 +1206,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[IntoDataFrameT] | Series: ... +) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload @@ -1265,7 +1266,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series: ... +) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series[IntoSeriesT]: ... @overload @@ -1277,7 +1278,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., -) -> Series: ... +) -> Series[IntoSeriesT]: ... @overload @@ -1337,7 +1338,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[Any] | LazyFrame[Any] | Series: ... +) -> DataFrame[Any] | LazyFrame[Any] | Series[IntoSeriesT]: ... @overload @@ -1349,7 +1350,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., -) -> Series: ... +) -> Series[IntoSeriesT]: ... @overload @@ -1385,7 +1386,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[IntoDataFrameT] | Series: ... +) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload @@ -1445,7 +1446,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series: ... +) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series[IntoSeriesT]: ... @overload @@ -1457,7 +1458,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., -) -> Series: ... +) -> Series[IntoSeriesT]: ... @overload @@ -1517,7 +1518,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[Any] | LazyFrame[Any] | Series: ... +) -> DataFrame[Any] | LazyFrame[Any] | Series[IntoSeriesT]: ... @overload @@ -1529,7 +1530,7 @@ def from_native( eager_or_interchange_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., -) -> Series: ... +) -> Series[IntoSeriesT]: ... @overload @@ -1557,8 +1558,8 @@ def from_native( ) -> Any: ... -def from_native( - native_object: IntoFrameT | IntoSeries | T, +def from_native( # type: ignore[misc] + native_object: IntoFrameT | IntoSeriesT | T, *, strict: bool | None = None, pass_through: bool | None = None, @@ -1566,7 +1567,7 @@ def from_native( eager_or_interchange_only: bool = False, series_only: bool = False, allow_series: bool | None = None, -) -> LazyFrame[IntoFrameT] | DataFrame[IntoFrameT] | Series | T: +) -> LazyFrame[IntoFrameT] | DataFrame[IntoFrameT] | Series[IntoSeriesT] | T: """Convert `native_object` to Narwhals Dataframe, Lazyframe, or Series. Arguments: @@ -1649,7 +1650,9 @@ def to_native( narwhals_object: LazyFrame[IntoFrameT], *, strict: Literal[True] = ... ) -> IntoFrameT: ... @overload -def to_native(narwhals_object: Series, *, strict: Literal[True] = ...) -> Any: ... +def to_native( + narwhals_object: Series[IntoSeriesT], *, strict: Literal[True] = ... +) -> IntoSeriesT: ... @overload def to_native(narwhals_object: Any, *, strict: bool) -> Any: ... @overload @@ -1661,17 +1664,22 @@ def to_native( narwhals_object: LazyFrame[IntoFrameT], *, pass_through: Literal[False] = ... ) -> IntoFrameT: ... @overload -def to_native(narwhals_object: Series, *, pass_through: Literal[False] = ...) -> Any: ... +def to_native( + narwhals_object: Series[IntoSeriesT], *, pass_through: Literal[False] = ... +) -> IntoSeriesT: ... @overload def to_native(narwhals_object: Any, *, pass_through: bool) -> Any: ... def to_native( - narwhals_object: DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series, + narwhals_object: DataFrame[IntoDataFrameT] + | LazyFrame[IntoFrameT] + | Series[IntoSeriesT] + | Any, *, strict: bool | None = None, pass_through: bool | None = None, -) -> IntoFrameT | Any: +) -> IntoDataFrameT | IntoFrameT | IntoSeriesT | Any: """Convert Narwhals object to native one. Arguments: @@ -2999,7 +3007,7 @@ def new_series( dtype: DType | type[DType] | None = None, *, native_namespace: ModuleType, -) -> Series: +) -> Series[Any]: """Instantiate Narwhals Series from iterable (e.g. list or array). Arguments: diff --git a/narwhals/translate.py b/narwhals/translate.py index e69c9a816..abb6ea8ac 100644 --- a/narwhals/translate.py +++ b/narwhals/translate.py @@ -66,7 +66,7 @@ def to_native( @overload def to_native( narwhals_object: Series[IntoSeriesT], *, pass_through: Literal[False] = ... -) -> Any: ... +) -> IntoSeriesT: ... @overload def to_native(narwhals_object: Any, *, pass_through: bool) -> Any: ... @@ -76,7 +76,7 @@ def to_native( *, strict: bool | None = None, pass_through: bool | None = None, -) -> IntoFrameT | Any: +) -> IntoFrameT | IntoSeriesT | Any: """Convert Narwhals object to native one. Arguments: @@ -310,8 +310,8 @@ def from_native( ) -> Any: ... -def from_native( # type: ignore[misc] - native_object: IntoFrameT | IntoSeriesT | T, +def from_native( + native_object: IntoFrameT | IntoSeriesT | Any, *, strict: bool | None = None, pass_through: bool | None = None, @@ -319,7 +319,7 @@ def from_native( # type: ignore[misc] eager_or_interchange_only: bool = False, series_only: bool = False, allow_series: bool | None = None, -) -> LazyFrame[IntoFrameT] | DataFrame[IntoFrameT] | Series[IntoSeriesT] | T: +) -> LazyFrame[IntoFrameT] | DataFrame[IntoFrameT] | Series[IntoSeriesT] | Any: """Convert `native_object` to Narwhals Dataframe, Lazyframe, or Series. Arguments: @@ -376,7 +376,7 @@ def from_native( # type: ignore[misc] strict, pass_through, pass_through_default=False, emit_deprecation_warning=True ) - return _from_native_impl( # type: ignore[no-any-return] + return _from_native_impl( native_object, pass_through=pass_through, eager_only=eager_only,