done
This commit is contained in:
481
lib/python3.11/site-packages/narwhals/typing.py
Normal file
481
lib/python3.11/site-packages/narwhals/typing.py
Normal file
@ -0,0 +1,481 @@
|
||||
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",
|
||||
]
|
Reference in New Issue
Block a user