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 import os 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. """ 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] FileSource: TypeAlias = "str | os.PathLike[str]" """Path to a file. Either a string or an object that implements [`__fspath__`], such as [`pathlib.Path`]. [`__fspath__`]: https://docs.python.org/3/library/os.html#os.PathLike [`pathlib.Path`]: https://docs.python.org/3/library/pathlib.html#pathlib.Path """ # 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", ]