diff --git a/narwhals/stable/v1/__init__.py b/narwhals/stable/v1/__init__.py index c43143a0ae..cececde65d 100644 --- a/narwhals/stable/v1/__init__.py +++ b/narwhals/stable/v1/__init__.py @@ -7,7 +7,6 @@ from typing import Iterable from typing import Literal from typing import Sequence -from typing import TypeVar from typing import overload from warnings import warn @@ -86,6 +85,7 @@ from typing import Mapping from typing_extensions import Self + from typing_extensions import TypeVar from narwhals.dtypes import DType from narwhals.functions import ArrowStreamExportable @@ -95,7 +95,13 @@ from narwhals.typing import _1DArray from narwhals.typing import _2DArray -T = TypeVar("T") + IntoSeriesT = TypeVar("IntoSeriesT", bound="IntoSeries", default=Any) + T = TypeVar("T", default=Any) +else: + from typing import TypeVar + + IntoSeriesT = TypeVar("IntoSeriesT", bound="IntoSeries") + T = TypeVar("T") class DataFrame(NwDataFrame[IntoDataFrameT]): @@ -127,7 +133,7 @@ class DataFrame(NwDataFrame[IntoDataFrameT]): # annotations are correct. @property - def _series(self: Self) -> type[Series]: + def _series(self: Self) -> type[Series[Any]]: return Series @property @@ -138,7 +144,7 @@ def _lazyframe(self: Self) -> type[LazyFrame[Any]]: def __getitem__( # type: ignore[overload-overlap] self: Self, item: str | tuple[slice | Sequence[int] | _1DArray, int | str], - ) -> Series: ... + ) -> Series[Any]: ... @overload def __getitem__( self: Self, @@ -192,16 +198,18 @@ def lazy( # 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: Self, *, as_series: Literal[True] = ...) -> dict[str, Series]: ... + def to_dict( + self: Self, *, as_series: Literal[True] = ... + ) -> dict[str, Series[Any]]: ... @overload def to_dict(self: Self, *, as_series: Literal[False]) -> dict[str, list[Any]]: ... @overload def to_dict( self: Self, *, as_series: bool - ) -> dict[str, Series] | dict[str, list[Any]]: ... + ) -> dict[str, Series[Any]] | dict[str, list[Any]]: ... def to_dict( self: 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: @@ -213,7 +221,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: @@ -221,7 +229,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: @@ -341,7 +349,7 @@ def tail(self, n: int = 5) -> Self: # pragma: no cover return super().tail(n) -class Series(NwSeries[Any]): +class Series(NwSeries[IntoSeriesT]): """Narwhals Series, backed by a native series. !!! warning @@ -1102,7 +1110,7 @@ def _stableify(obj: NwDataFrame[IntoFrameT]) -> DataFrame[IntoFrameT]: ... @overload def _stableify(obj: NwLazyFrame[IntoFrameT]) -> LazyFrame[IntoFrameT]: ... @overload -def _stableify(obj: NwSeries[Any]) -> Series: ... +def _stableify(obj: NwSeries[IntoSeriesT]) -> Series[IntoSeriesT]: ... @overload def _stableify(obj: NwExpr) -> Expr: ... @overload @@ -1110,8 +1118,12 @@ def _stableify(obj: Any) -> Any: ... def _stableify( - obj: NwDataFrame[IntoFrameT] | NwLazyFrame[IntoFrameT] | NwSeries[Any] | NwExpr | Any, -) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series | Expr | Any: + obj: NwDataFrame[IntoFrameT] + | NwLazyFrame[IntoFrameT] + | NwSeries[IntoSeriesT] + | NwExpr + | Any, +) -> DataFrame[IntoFrameT] | LazyFrame[IntoFrameT] | Series[IntoSeriesT] | Expr | Any: if isinstance(obj, NwDataFrame): return DataFrame( obj._compliant_frame._change_version(Version.V1), @@ -1134,26 +1146,26 @@ def _stableify( @overload def from_native( - native_object: IntoDataFrameT | IntoSeries, + native_object: IntoDataFrameT | IntoSeriesT, *, strict: Literal[False], eager_only: Literal[False] = ..., eager_or_interchange_only: Literal[True], series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[IntoDataFrameT] | Series: ... +) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload def from_native( - native_object: IntoDataFrameT | IntoSeries, + native_object: IntoDataFrameT | IntoSeriesT, *, strict: Literal[False], eager_only: Literal[True], eager_or_interchange_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[IntoDataFrameT] | Series: ... +) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload @@ -1206,26 +1218,26 @@ def from_native( @overload def from_native( - native_object: IntoFrameT | IntoSeries, + native_object: IntoFrameT | IntoSeriesT, *, strict: Literal[False], eager_only: Literal[False] = ..., 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 def from_native( - native_object: IntoSeries, + native_object: IntoSeriesT, *, strict: Literal[False], eager_only: Literal[False] = ..., eager_or_interchange_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., -) -> Series: ... +) -> Series[IntoSeriesT]: ... @overload @@ -1285,19 +1297,19 @@ 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[Any]: ... @overload def from_native( - native_object: IntoSeries | Any, # remain `Any` for downstream compatibility + native_object: IntoSeriesT, *, strict: Literal[True] = ..., eager_only: Literal[False] = ..., eager_or_interchange_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., -) -> Series: ... +) -> Series[IntoSeriesT]: ... @overload @@ -1326,14 +1338,14 @@ def from_native( @overload def from_native( - native_object: IntoDataFrameT | IntoSeries, + native_object: IntoDataFrameT | IntoSeriesT, *, pass_through: Literal[True], eager_only: Literal[True], eager_or_interchange_only: Literal[False] = ..., series_only: Literal[False] = ..., allow_series: Literal[True], -) -> DataFrame[IntoDataFrameT] | Series: ... +) -> DataFrame[IntoDataFrameT] | Series[IntoSeriesT]: ... @overload @@ -1386,26 +1398,26 @@ def from_native( @overload def from_native( - native_object: IntoFrameT | IntoSeries, + native_object: IntoFrameT | IntoSeriesT, *, pass_through: Literal[True], eager_only: Literal[False] = ..., 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 def from_native( - native_object: IntoSeries, + native_object: IntoSeriesT, *, pass_through: Literal[True], eager_only: Literal[False] = ..., eager_or_interchange_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., -) -> Series: ... +) -> Series[IntoSeriesT]: ... @overload @@ -1465,19 +1477,19 @@ 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[Any]: ... @overload def from_native( - native_object: IntoSeries, + native_object: IntoSeriesT, *, pass_through: Literal[False] = ..., eager_only: Literal[False] = ..., eager_or_interchange_only: Literal[False] = ..., series_only: Literal[True], allow_series: None = ..., -) -> Series: ... +) -> Series[IntoSeriesT]: ... @overload @@ -1506,7 +1518,7 @@ def from_native( def from_native( - native_object: IntoFrameT | IntoFrame | IntoSeries | T, + native_object: IntoFrameT | IntoFrame | IntoSeriesT | IntoSeries | T, *, strict: bool | None = None, pass_through: bool | None = None, @@ -1514,7 +1526,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: @@ -1595,7 +1607,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 @@ -1607,17 +1621,21 @@ 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[IntoDataFrameT] | LazyFrame[IntoFrameT] | Series, + narwhals_object: DataFrame[IntoDataFrameT] + | LazyFrame[IntoFrameT] + | Series[IntoSeriesT], *, strict: bool | None = None, pass_through: bool | None = None, -) -> IntoFrameT | Any: +) -> IntoFrameT | IntoSeriesT | Any: """Convert Narwhals object to native one. Arguments: @@ -2124,7 +2142,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/stable/v1/typing.py b/narwhals/stable/v1/typing.py index fa7cd67fc7..94df953f3a 100644 --- a/narwhals/stable/v1/typing.py +++ b/narwhals/stable/v1/typing.py @@ -1,215 +1,215 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING -from typing import Any -from typing import Protocol -from typing import TypeVar -from typing import Union - -if TYPE_CHECKING: - import sys - - if sys.version_info >= (3, 10): - from typing import TypeAlias - else: - from typing_extensions import TypeAlias - - from narwhals.stable.v1 import DataFrame - from narwhals.stable.v1 import Expr - from narwhals.stable.v1 import LazyFrame - from narwhals.stable.v1 import Series - from narwhals.stable.v1 import dtypes - - # All dataframes supported by Narwhals have a - # `columns` property. Their similarities don't extend - # _that_ much further unfortunately... - class NativeFrame(Protocol): - @property - def columns(self) -> Any: ... - - def join(self, *args: Any, **kwargs: Any) -> Any: ... - - class NativeSeries(Protocol): - def __len__(self) -> int: ... - - class DataFrameLike(Protocol): - def __dataframe__(self, *args: Any, **kwargs: Any) -> Any: ... - - -IntoExpr: TypeAlias = Union["Expr", str, "Series"] -"""Anything which can be converted to an expression. - -Use this to mean "either a Narwhals expression, or something -which can be converted into one". For example, `exprs` in `DataFrame.select` is -typed to accept `IntoExpr`, as it can either accept a `nw.Expr` -(e.g. `df.select(nw.col('a'))`) or a string which will be interpreted as a -`nw.Expr`, e.g. `df.select('a')`. -""" - -IntoDataFrame: TypeAlias = Union["NativeFrame", "DataFrame[Any]", "DataFrameLike"] -"""Anything which can be converted to a Narwhals DataFrame. - -Use this if your function accepts a narwhalifiable object but doesn't care about its backend. - -Examples: - >>> import narwhals as nw - >>> from narwhals.typing import IntoDataFrame - >>> def agnostic_shape(df_native: IntoDataFrame) -> tuple[int, int]: - ... df = nw.from_native(df_native, eager_only=True) - ... return df.shape -""" - -IntoFrame: TypeAlias = Union[ - "NativeFrame", "DataFrame[Any]", "LazyFrame[Any]", "DataFrameLike" -] -"""Anything which can be converted to a Narwhals DataFrame or LazyFrame. - -Use this if your function can accept an object which can be converted to either -`nw.DataFrame` or `nw.LazyFrame` and it doesn't care about its backend. - -Examples: - >>> import narwhals as nw - >>> from narwhals.typing import IntoFrame - >>> def agnostic_columns(df_native: IntoFrame) -> list[str]: - ... df = nw.from_native(df_native) - ... return df.collect_schema().names() -""" - -Frame: TypeAlias = Union["DataFrame[Any]", "LazyFrame[Any]"] -"""Narwhals DataFrame or Narwhals LazyFrame. - -Use this if your function can work with either and your function doesn't care -about its backend. - -Examples: - >>> import narwhals as nw - >>> from narwhals.typing import Frame - >>> @nw.narwhalify - ... def agnostic_columns(df: Frame) -> list[str]: - ... return df.columns -""" - -IntoSeries: TypeAlias = Union["Series", "NativeSeries"] -"""Anything which can be converted to a Narwhals Series. - -Use this if your function can accept an object which can be converted to `nw.Series` -and it doesn't care about its backend. - -Examples: - >>> from typing import Any - >>> import narwhals as nw - >>> from narwhals.typing import IntoSeries - >>> def agnostic_to_list(s_native: IntoSeries) -> list[Any]: - ... s = nw.from_native(s_native) - ... return s.to_list() -""" - -IntoFrameT = TypeVar("IntoFrameT", bound="IntoFrame") -"""TypeVar bound to object convertible to Narwhals DataFrame or Narwhals LazyFrame. - -Use this if your function accepts an object which is convertible to `nw.DataFrame` -or `nw.LazyFrame` and returns an object of the same type. - -Examples: - >>> import narwhals as nw - >>> from narwhals.typing import IntoFrameT - >>> def agnostic_func(df_native: IntoFrameT) -> IntoFrameT: - ... df = nw.from_native(df_native) - ... return df.with_columns(c=nw.col("a") + 1).to_native() -""" - -IntoDataFrameT = TypeVar("IntoDataFrameT", bound="IntoDataFrame") -"""TypeVar bound to object convertible to Narwhals DataFrame. - -Use this if your function accepts an object which can be converted to `nw.DataFrame` -and returns an object of the same class. - -Examples: - >>> import narwhals as nw - >>> from narwhals.typing import IntoDataFrameT - >>> def agnostic_func(df_native: IntoDataFrameT) -> IntoDataFrameT: - ... df = nw.from_native(df_native, eager_only=True) - ... return df.with_columns(c=df["a"] + 1).to_native() -""" - -FrameT = TypeVar("FrameT", bound="Frame") -"""TypeVar bound to Narwhals DataFrame or Narwhals LazyFrame. - -Use this if your function accepts either `nw.DataFrame` or `nw.LazyFrame` and returns -an object of the same kind. - -Examples: - >>> import narwhals as nw - >>> from narwhals.typing import FrameT - >>> @nw.narwhalify - ... def agnostic_func(df: FrameT) -> FrameT: - ... return df.with_columns(c=nw.col("a") + 1) -""" - -DataFrameT = TypeVar("DataFrameT", bound="DataFrame[Any]") -"""TypeVar bound to Narwhals DataFrame. - -Use this if your function can accept a Narwhals DataFrame and returns a Narwhals -DataFrame backed by the same backend. - -Examples: - >>> import narwhals as nw - >>> from narwhals.typing import DataFrameT - >>> @nw.narwhalify - >>> def func(df: DataFrameT) -> DataFrameT: - ... return df.with_columns(c=df["a"] + 1) -""" - -IntoSeriesT = TypeVar("IntoSeriesT", bound="IntoSeries") -"""TypeVar bound to object convertible to Narwhals Series. - -Use this if your function accepts an object which can be converted to `nw.Series` -and returns an object of the same class. - -Examples: - >>> import narwhals as nw - >>> from narwhals.typing import IntoSeriesT - >>> def agnostic_abs(s_native: IntoSeriesT) -> IntoSeriesT: - ... s = nw.from_native(s_native, series_only=True) - ... return s.abs().to_native() -""" - - -class DTypes: - Int64: type[dtypes.Int64] - Int32: type[dtypes.Int32] - Int16: type[dtypes.Int16] - Int8: type[dtypes.Int8] - UInt64: type[dtypes.UInt64] - UInt32: type[dtypes.UInt32] - UInt16: type[dtypes.UInt16] - UInt8: type[dtypes.UInt8] - Float64: type[dtypes.Float64] - Float32: type[dtypes.Float32] - String: type[dtypes.String] - Boolean: type[dtypes.Boolean] - Object: type[dtypes.Object] - Categorical: type[dtypes.Categorical] - Enum: type[dtypes.Enum] - Datetime: type[dtypes.Datetime] - Duration: type[dtypes.Duration] - Date: type[dtypes.Date] - Field: type[dtypes.Field] - Struct: type[dtypes.Struct] - List: type[dtypes.List] - Array: type[dtypes.Array] - Unknown: type[dtypes.Unknown] - - -__all__ = [ - "DataFrameT", - "Frame", - "FrameT", - "IntoDataFrame", - "IntoDataFrameT", - "IntoExpr", - "IntoFrame", - "IntoFrameT", - "IntoSeries", - "IntoSeriesT", -] +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing import Any +from typing import Protocol +from typing import TypeVar +from typing import Union + +if TYPE_CHECKING: + import sys + + if sys.version_info >= (3, 10): + from typing import TypeAlias + else: + from typing_extensions import TypeAlias + + from narwhals.stable.v1 import DataFrame + from narwhals.stable.v1 import Expr + from narwhals.stable.v1 import LazyFrame + from narwhals.stable.v1 import Series + from narwhals.stable.v1 import dtypes + + # All dataframes supported by Narwhals have a + # `columns` property. Their similarities don't extend + # _that_ much further unfortunately... + class NativeFrame(Protocol): + @property + def columns(self) -> Any: ... + + def join(self, *args: Any, **kwargs: Any) -> Any: ... + + class NativeSeries(Protocol): + def __len__(self) -> int: ... + + class DataFrameLike(Protocol): + def __dataframe__(self, *args: Any, **kwargs: Any) -> Any: ... + + +IntoExpr: TypeAlias = Union["Expr", str, "Series[Any]"] +"""Anything which can be converted to an expression. + +Use this to mean "either a Narwhals expression, or something +which can be converted into one". For example, `exprs` in `DataFrame.select` is +typed to accept `IntoExpr`, as it can either accept a `nw.Expr` +(e.g. `df.select(nw.col('a'))`) or a string which will be interpreted as a +`nw.Expr`, e.g. `df.select('a')`. +""" + +IntoDataFrame: TypeAlias = Union["NativeFrame", "DataFrame[Any]", "DataFrameLike"] +"""Anything which can be converted to a Narwhals DataFrame. + +Use this if your function accepts a narwhalifiable object but doesn't care about its backend. + +Examples: + >>> import narwhals as nw + >>> from narwhals.typing import IntoDataFrame + >>> def agnostic_shape(df_native: IntoDataFrame) -> tuple[int, int]: + ... df = nw.from_native(df_native, eager_only=True) + ... return df.shape +""" + +IntoFrame: TypeAlias = Union[ + "NativeFrame", "DataFrame[Any]", "LazyFrame[Any]", "DataFrameLike" +] +"""Anything which can be converted to a Narwhals DataFrame or LazyFrame. + +Use this if your function can accept an object which can be converted to either +`nw.DataFrame` or `nw.LazyFrame` and it doesn't care about its backend. + +Examples: + >>> import narwhals as nw + >>> from narwhals.typing import IntoFrame + >>> def agnostic_columns(df_native: IntoFrame) -> list[str]: + ... df = nw.from_native(df_native) + ... return df.collect_schema().names() +""" + +Frame: TypeAlias = Union["DataFrame[Any]", "LazyFrame[Any]"] +"""Narwhals DataFrame or Narwhals LazyFrame. + +Use this if your function can work with either and your function doesn't care +about its backend. + +Examples: + >>> import narwhals as nw + >>> from narwhals.typing import Frame + >>> @nw.narwhalify + ... def agnostic_columns(df: Frame) -> list[str]: + ... return df.columns +""" + +IntoSeries: TypeAlias = Union["Series[Any]", "NativeSeries"] +"""Anything which can be converted to a Narwhals Series. + +Use this if your function can accept an object which can be converted to `nw.Series` +and it doesn't care about its backend. + +Examples: + >>> from typing import Any + >>> import narwhals as nw + >>> from narwhals.typing import IntoSeries + >>> def agnostic_to_list(s_native: IntoSeries) -> list[Any]: + ... s = nw.from_native(s_native) + ... return s.to_list() +""" + +IntoFrameT = TypeVar("IntoFrameT", bound="IntoFrame") +"""TypeVar bound to object convertible to Narwhals DataFrame or Narwhals LazyFrame. + +Use this if your function accepts an object which is convertible to `nw.DataFrame` +or `nw.LazyFrame` and returns an object of the same type. + +Examples: + >>> import narwhals as nw + >>> from narwhals.typing import IntoFrameT + >>> def agnostic_func(df_native: IntoFrameT) -> IntoFrameT: + ... df = nw.from_native(df_native) + ... return df.with_columns(c=nw.col("a") + 1).to_native() +""" + +IntoDataFrameT = TypeVar("IntoDataFrameT", bound="IntoDataFrame") +"""TypeVar bound to object convertible to Narwhals DataFrame. + +Use this if your function accepts an object which can be converted to `nw.DataFrame` +and returns an object of the same class. + +Examples: + >>> import narwhals as nw + >>> from narwhals.typing import IntoDataFrameT + >>> def agnostic_func(df_native: IntoDataFrameT) -> IntoDataFrameT: + ... df = nw.from_native(df_native, eager_only=True) + ... return df.with_columns(c=df["a"] + 1).to_native() +""" + +FrameT = TypeVar("FrameT", bound="Frame") +"""TypeVar bound to Narwhals DataFrame or Narwhals LazyFrame. + +Use this if your function accepts either `nw.DataFrame` or `nw.LazyFrame` and returns +an object of the same kind. + +Examples: + >>> import narwhals as nw + >>> from narwhals.typing import FrameT + >>> @nw.narwhalify + ... def agnostic_func(df: FrameT) -> FrameT: + ... return df.with_columns(c=nw.col("a") + 1) +""" + +DataFrameT = TypeVar("DataFrameT", bound="DataFrame[Any]") +"""TypeVar bound to Narwhals DataFrame. + +Use this if your function can accept a Narwhals DataFrame and returns a Narwhals +DataFrame backed by the same backend. + +Examples: + >>> import narwhals as nw + >>> from narwhals.typing import DataFrameT + >>> @nw.narwhalify + >>> def func(df: DataFrameT) -> DataFrameT: + ... return df.with_columns(c=df["a"] + 1) +""" + +IntoSeriesT = TypeVar("IntoSeriesT", bound="IntoSeries") +"""TypeVar bound to object convertible to Narwhals Series. + +Use this if your function accepts an object which can be converted to `nw.Series` +and returns an object of the same class. + +Examples: + >>> import narwhals as nw + >>> from narwhals.typing import IntoSeriesT + >>> def agnostic_abs(s_native: IntoSeriesT) -> IntoSeriesT: + ... s = nw.from_native(s_native, series_only=True) + ... return s.abs().to_native() +""" + + +class DTypes: + Int64: type[dtypes.Int64] + Int32: type[dtypes.Int32] + Int16: type[dtypes.Int16] + Int8: type[dtypes.Int8] + UInt64: type[dtypes.UInt64] + UInt32: type[dtypes.UInt32] + UInt16: type[dtypes.UInt16] + UInt8: type[dtypes.UInt8] + Float64: type[dtypes.Float64] + Float32: type[dtypes.Float32] + String: type[dtypes.String] + Boolean: type[dtypes.Boolean] + Object: type[dtypes.Object] + Categorical: type[dtypes.Categorical] + Enum: type[dtypes.Enum] + Datetime: type[dtypes.Datetime] + Duration: type[dtypes.Duration] + Date: type[dtypes.Date] + Field: type[dtypes.Field] + Struct: type[dtypes.Struct] + List: type[dtypes.List] + Array: type[dtypes.Array] + Unknown: type[dtypes.Unknown] + + +__all__ = [ + "DataFrameT", + "Frame", + "FrameT", + "IntoDataFrame", + "IntoDataFrameT", + "IntoExpr", + "IntoFrame", + "IntoFrameT", + "IntoSeries", + "IntoSeriesT", +] diff --git a/tests/expr_and_series/dt/timestamp_test.py b/tests/expr_and_series/dt/timestamp_test.py index 448091cad4..ed22a16dbc 100644 --- a/tests/expr_and_series/dt/timestamp_test.py +++ b/tests/expr_and_series/dt/timestamp_test.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import datetime +from typing import TYPE_CHECKING from typing import Literal import hypothesis.strategies as st @@ -18,6 +19,9 @@ from tests.utils import assert_equal_data from tests.utils import is_windows +if TYPE_CHECKING: + from narwhals.typing import IntoSeriesT + data = { "a": [ datetime(2021, 3, 1, 12, 34, 56, 49000), @@ -221,7 +225,7 @@ def test_timestamp_hypothesis( import polars as pl @nw.narwhalify - def func(s: nw.Series) -> nw.Series: + def func(s: nw.Series[IntoSeriesT]) -> nw.Series[IntoSeriesT]: return s.dt.timestamp(time_unit) result_pl = func(pl.Series([inputs], dtype=pl.Datetime(starting_time_unit))) diff --git a/tests/translate/from_native_test.py b/tests/translate/from_native_test.py index 9ad0ec5b8a..7b48c8a79f 100644 --- a/tests/translate/from_native_test.py +++ b/tests/translate/from_native_test.py @@ -3,6 +3,9 @@ from contextlib import nullcontext as does_not_raise from typing import TYPE_CHECKING from typing import Any +from typing import Iterable +from typing import Literal +from typing import cast import numpy as np import pandas as pd @@ -241,7 +244,7 @@ def test_series_only_sqlframe() -> None: # pragma: no cover df = session.createDataFrame([*zip(*data.values())], schema=[*data.keys()]) with pytest.raises(TypeError, match="Cannot only use `series_only`"): - nw.from_native(df, series_only=True) + nw.from_native(df, series_only=True) # pyright: ignore[reportArgumentType, reportCallIssue] @pytest.mark.parametrize( @@ -287,3 +290,18 @@ def __dataframe__(self) -> None: # pragma: no cover mockdf = MockDf() result = nw.from_native(mockdf, eager_only=True, strict=False) assert result is mockdf + + +def test_from_native_altair_array_like() -> None: + obj: list[int] = [1, 2, 3, 4] + array_like = cast("Iterable[Any]", obj) + not_array_like: Literal[1] = 1 + + with pytest.raises(TypeError, match="got.+list"): + false_positive_native_series = nw.from_native(obj, series_only=True) # noqa: F841 + + with pytest.raises(TypeError, match="got.+list"): + true_negative_iterable = nw.from_native(array_like, series_only=True) # type: ignore[call-overload] # noqa: F841 + + with pytest.raises(TypeError, match="got.+int"): + true_negative_not_native_series = nw.from_native(not_array_like, series_only=True) # type: ignore[call-overload] # noqa: F841