183 lines
5.8 KiB
Python
183 lines
5.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any, Protocol, TypeVar, Union
|
|
|
|
if TYPE_CHECKING:
|
|
import sys
|
|
|
|
from narwhals.stable.v1 import DataFrame, LazyFrame
|
|
|
|
if sys.version_info >= (3, 10):
|
|
from typing import TypeAlias
|
|
else:
|
|
from typing_extensions import TypeAlias
|
|
|
|
from narwhals.stable.v1 import Expr, Series
|
|
|
|
# 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", "DataFrame[Any]", "LazyFrame[Any]")
|
|
"""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()
|
|
"""
|
|
|
|
|
|
__all__ = [
|
|
"DataFrameT",
|
|
"Frame",
|
|
"FrameT",
|
|
"IntoDataFrame",
|
|
"IntoDataFrameT",
|
|
"IntoExpr",
|
|
"IntoFrame",
|
|
"IntoFrameT",
|
|
"IntoSeries",
|
|
"IntoSeriesT",
|
|
]
|