team-10/env/Lib/site-packages/narwhals/stable/v1/typing.py
2025-08-02 07:34:44 +02:00

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",
]