Files
dash-api/lib/python3.11/site-packages/narwhals/typing.py

482 lines
17 KiB
Python
Raw Normal View History

2025-09-07 22:09:54 +02:00
from __future__ import annotations
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, Union
from narwhals._compliant import CompliantDataFrame, CompliantLazyFrame, CompliantSeries
from narwhals._typing import Backend, EagerAllowed, IntoBackend, LazyAllowed
if TYPE_CHECKING:
import datetime as dt
from collections.abc import Iterable, Sequence, Sized
from decimal import Decimal
from types import ModuleType
import numpy as np
import pandas as pd
import polars as pl
import pyarrow as pa
from typing_extensions import TypeAlias
from narwhals import dtypes
from narwhals._namespace import _NativeIbis
from narwhals.dataframe import DataFrame, LazyFrame
from narwhals.expr import Expr
from narwhals.schema import Schema
from narwhals.series import 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 NativeDataFrame(Sized, NativeFrame, Protocol): ...
class NativeLazyFrame(NativeFrame, Protocol):
def explain(self, *args: Any, **kwargs: Any) -> Any: ...
class NativeSeries(Sized, Iterable[Any], Protocol):
def filter(self, *args: Any, **kwargs: Any) -> Any: ...
class SupportsNativeNamespace(Protocol):
def __native_namespace__(self) -> ModuleType: ...
# ruff: noqa: N802
class DTypes(Protocol):
@property
def Decimal(self) -> type[dtypes.Decimal]: ...
@property
def Int128(self) -> type[dtypes.Int128]: ...
@property
def Int64(self) -> type[dtypes.Int64]: ...
@property
def Int32(self) -> type[dtypes.Int32]: ...
@property
def Int16(self) -> type[dtypes.Int16]: ...
@property
def Int8(self) -> type[dtypes.Int8]: ...
@property
def UInt128(self) -> type[dtypes.UInt128]: ...
@property
def UInt64(self) -> type[dtypes.UInt64]: ...
@property
def UInt32(self) -> type[dtypes.UInt32]: ...
@property
def UInt16(self) -> type[dtypes.UInt16]: ...
@property
def UInt8(self) -> type[dtypes.UInt8]: ...
@property
def Float64(self) -> type[dtypes.Float64]: ...
@property
def Float32(self) -> type[dtypes.Float32]: ...
@property
def String(self) -> type[dtypes.String]: ...
@property
def Boolean(self) -> type[dtypes.Boolean]: ...
@property
def Object(self) -> type[dtypes.Object]: ...
@property
def Categorical(self) -> type[dtypes.Categorical]: ...
@property
def Enum(self) -> type[dtypes.Enum]: ...
@property
def Datetime(self) -> type[dtypes.Datetime]: ...
@property
def Duration(self) -> type[dtypes.Duration]: ...
@property
def Date(self) -> type[dtypes.Date]: ...
@property
def Field(self) -> type[dtypes.Field]: ...
@property
def Struct(self) -> type[dtypes.Struct]: ...
@property
def List(self) -> type[dtypes.List]: ...
@property
def Array(self) -> type[dtypes.Array]: ...
@property
def Unknown(self) -> type[dtypes.Unknown]: ...
@property
def Time(self) -> type[dtypes.Time]: ...
@property
def Binary(self) -> type[dtypes.Binary]: ...
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 = "NativeDataFrame"
"""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
"""
IntoLazyFrame: TypeAlias = Union["NativeLazyFrame", "_NativeIbis"]
IntoFrame: TypeAlias = Union["IntoDataFrame", "IntoLazyFrame"]
"""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 = "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()
"""
IntoLazyFrameT = TypeVar("IntoLazyFrameT", bound="IntoLazyFrame")
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)
"""
LazyFrameT = TypeVar("LazyFrameT", bound="LazyFrame[Any]")
SeriesT = TypeVar("SeriesT", bound="Series[Any]")
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()
"""
DTypeBackend: TypeAlias = 'Literal["pyarrow", "numpy_nullable"] | None'
SizeUnit: TypeAlias = Literal[
"b",
"kb",
"mb",
"gb",
"tb",
"bytes",
"kilobytes",
"megabytes",
"gigabytes",
"terabytes",
]
TimeUnit: TypeAlias = Literal["ns", "us", "ms", "s"]
AsofJoinStrategy: TypeAlias = Literal["backward", "forward", "nearest"]
"""Join strategy.
- *"backward"*: Selects the last row in the right DataFrame whose `on` key
is less than or equal to the left's key.
- *"forward"*: Selects the first row in the right DataFrame whose `on` key
is greater than or equal to the left's key.
- *"nearest"*: Search selects the last row in the right DataFrame whose value
is nearest to the left's key.
"""
ClosedInterval: TypeAlias = Literal["left", "right", "none", "both"]
"""Define which sides of the interval are closed (inclusive)."""
ConcatMethod: TypeAlias = Literal["horizontal", "vertical", "diagonal"]
"""Concatenating strategy.
- *"vertical"*: Concatenate vertically. Column names must match.
- *"horizontal"*: Concatenate horizontally. If lengths don't match, then
missing rows are filled with null values.
- *"diagonal"*: Finds a union between the column schemas and fills missing
column values with null.
"""
FillNullStrategy: TypeAlias = Literal["forward", "backward"]
"""Strategy used to fill null values."""
JoinStrategy: TypeAlias = Literal["inner", "left", "full", "cross", "semi", "anti"]
"""Join strategy.
- *"inner"*: Returns rows that have matching values in both tables.
- *"left"*: Returns all rows from the left table, and the matched rows from
the right table.
- *"full"*: Returns all rows in both dataframes, with the `suffix` appended to
the right join keys.
- *"cross"*: Returns the Cartesian product of rows from both tables.
- *"semi"*: Filter rows that have a match in the right table.
- *"anti"*: Filter rows that do not have a match in the right table.
"""
PivotAgg: TypeAlias = Literal[
"min", "max", "first", "last", "sum", "mean", "median", "len"
]
"""A predefined aggregate function string."""
RankMethod: TypeAlias = Literal["average", "min", "max", "dense", "ordinal"]
"""The method used to assign ranks to tied elements.
- *"average"*: The average of the ranks that would have been assigned to
all the tied values is assigned to each value.
- *"min"*: The minimum of the ranks that would have been assigned to all
the tied values is assigned to each value. (This is also referred to
as "competition" ranking.)
- *"max"*: The maximum of the ranks that would have been assigned to all
the tied values is assigned to each value.
- *"dense"*: Like "min", but the rank of the next highest element is
assigned the rank immediately after those assigned to the tied elements.
- *"ordinal"*: All values are given a distinct rank, corresponding to the
order that the values occur in the Series.
"""
RollingInterpolationMethod: TypeAlias = Literal[
"nearest", "higher", "lower", "midpoint", "linear"
]
"""Interpolation method."""
UniqueKeepStrategy: TypeAlias = Literal["any", "first", "last", "none"]
"""Which of the duplicate rows to keep.
- *"any"*: Does not give any guarantee of which row is kept.
This allows more optimizations.
- *"none"*: Don't keep duplicate rows.
- *"first"*: Keep first unique row.
- *"last"*: Keep last unique row.
"""
LazyUniqueKeepStrategy: TypeAlias = Literal["any", "none"]
"""Which of the duplicate rows to keep.
- *"any"*: Does not give any guarantee of which row is kept.
- *"none"*: Don't keep duplicate rows.
"""
ModeKeepStrategy: TypeAlias = Literal["any", "all"]
"""Which of the mode's to keep.
- *"any"*: Does not give any guarantee of which mode is kept.
- *"all"*: Keeps all the mode's.
"""
_ShapeT = TypeVar("_ShapeT", bound="tuple[int, ...]")
_NDArray: TypeAlias = "np.ndarray[_ShapeT, Any]"
_1DArray: TypeAlias = "_NDArray[tuple[int]]"
_1DArrayInt: TypeAlias = "np.ndarray[tuple[int], np.dtype[np.integer[Any]]]"
_2DArray: TypeAlias = "_NDArray[tuple[int, int]]" # noqa: PYI047
_AnyDArray: TypeAlias = "_NDArray[tuple[int, ...]]" # noqa: PYI047
_NumpyScalar: TypeAlias = "np.generic[Any]"
Into1DArray: TypeAlias = "_1DArray | _NumpyScalar"
"""A 1-dimensional `numpy.ndarray` or scalar that can be converted into one."""
PandasLikeDType: TypeAlias = "pd.api.extensions.ExtensionDtype | np.dtype[Any]"
NumericLiteral: TypeAlias = "int | float | Decimal"
TemporalLiteral: TypeAlias = "dt.date | dt.datetime | dt.time | dt.timedelta"
NonNestedLiteral: TypeAlias = (
"NumericLiteral | TemporalLiteral | str | bool | bytes | None"
)
PythonLiteral: TypeAlias = "NonNestedLiteral | list[Any] | tuple[Any, ...]"
NonNestedDType: TypeAlias = "dtypes.NumericType | dtypes.TemporalType | dtypes.String | dtypes.Boolean | dtypes.Binary | dtypes.Categorical | dtypes.Unknown | dtypes.Object"
"""Any Narwhals DType that does not have required arguments."""
IntoDType: TypeAlias = "dtypes.DType | type[NonNestedDType]"
"""Anything that can be converted into a Narwhals DType.
Examples:
>>> import polars as pl
>>> import narwhals as nw
>>> df_native = pl.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
>>> df = nw.from_native(df_native)
>>> df.select(
... nw.col("a").cast(nw.Int32),
... nw.col("b").cast(nw.String()).str.split(".").cast(nw.List(nw.Int8)),
... )
|Narwhals DataFrame|
|------------------|
|shape: (3, 2) |
||
| a b |
| --- --- |
| i32 list[i8] |
||
| 1 [4, 0] |
| 2 [5, 0] |
| 3 [6, 0] |
||
"""
# TODO @dangotbanned: fix this?
# Constructor allows tuples, but we don't support that *everywhere* yet
IntoSchema: TypeAlias = "Mapping[str, dtypes.DType] | Schema"
"""Anything that can be converted into a Narwhals Schema.
Defined by column names and their associated *instantiated* Narwhals DType.
Examples:
>>> import narwhals as nw
>>> import pyarrow as pa
>>> data = {"a": [1, 2, 3], "b": [None, "hi", "howdy"], "c": [2.1, 2.0, None]}
>>> nw.DataFrame.from_dict(
... data,
... schema={"a": nw.UInt8(), "b": nw.String(), "c": nw.Float32()},
... backend="pyarrow",
... )
| Narwhals DataFrame |
|------------------------|
|pyarrow.Table |
|a: uint8 |
|b: string |
|c: float |
|---- |
|a: [[1,2,3]] |
|b: [[null,"hi","howdy"]]|
|c: [[2.1,2,null]] |
"""
IntoArrowSchema: TypeAlias = "pa.Schema | Mapping[str, pa.DataType]"
IntoPolarsSchema: TypeAlias = "pl.Schema | Mapping[str, pl.DataType]"
IntoPandasSchema: TypeAlias = Mapping[str, PandasLikeDType]
# Annotations for `__getitem__` methods
_T = TypeVar("_T")
_Slice: TypeAlias = "slice[_T, Any, Any] | slice[Any, _T, Any] | slice[None, None, _T]"
_SliceNone: TypeAlias = "slice[None, None, None]"
# Index/column positions
SingleIndexSelector: TypeAlias = int
_SliceIndex: TypeAlias = "_Slice[int] | _SliceNone"
"""E.g. `[1:]` or `[:3]` or `[::2]`."""
SizedMultiIndexSelector: TypeAlias = "Sequence[int] | _T | _1DArrayInt"
MultiIndexSelector: TypeAlias = "_SliceIndex | SizedMultiIndexSelector[_T]"
# Labels/column names
SingleNameSelector: TypeAlias = str
_SliceName: TypeAlias = "_Slice[str] | _SliceNone"
SizedMultiNameSelector: TypeAlias = "Sequence[str] | _T | _1DArray"
MultiNameSelector: TypeAlias = "_SliceName | SizedMultiNameSelector[_T]"
# Mixed selectors
SingleColSelector: TypeAlias = "SingleIndexSelector | SingleNameSelector"
MultiColSelector: TypeAlias = "MultiIndexSelector[_T] | MultiNameSelector[_T]"
__all__ = [
"Backend",
"CompliantDataFrame",
"CompliantLazyFrame",
"CompliantSeries",
"DataFrameT",
"EagerAllowed",
"Frame",
"FrameT",
"IntoBackend",
"IntoDataFrame",
"IntoDataFrameT",
"IntoExpr",
"IntoFrame",
"IntoFrameT",
"IntoSeries",
"IntoSeriesT",
"LazyAllowed",
]