Skip to content

Commit

Permalink
refactor(typing): Add __getitem__ selector aliases
Browse files Browse the repository at this point in the history
  • Loading branch information
dangotbanned committed Feb 15, 2025
1 parent f79a9ea commit 8d8895b
Showing 1 changed file with 50 additions and 1 deletion.
51 changes: 50 additions & 1 deletion narwhals/typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -268,11 +268,60 @@ def __native_namespace__(self) -> ModuleType: ...

_ShapeT = TypeVar("_ShapeT", bound="tuple[int, ...]")
_NDArray: TypeAlias = "np.ndarray[_ShapeT, Any]"
_1DArray: TypeAlias = "_NDArray[tuple[int]]" # noqa: PYI042, PYI047
_1DArray: TypeAlias = "_NDArray[tuple[int]]" # noqa: PYI042
_2DArray: TypeAlias = "_NDArray[tuple[int, int]]" # noqa: PYI042, PYI047
_AnyDArray: TypeAlias = "_NDArray[tuple[int, ...]]" # noqa: PYI047


# Annotations for `__getitem__` methods
_Slice: TypeAlias = "slice[Any, Any, Any]"

SingleIndexSelector: TypeAlias = int
SingleNameSelector: TypeAlias = str

MultiIndexSelector: TypeAlias = Union[
_Slice,
Sequence[int],
_1DArray,
# IntoSeriesT, Would need support like (#2013)
]
MultiNameSelector: TypeAlias = Union[
_Slice,
Sequence[str],
_1DArray,
# IntoSeriesT, Would need support like (#2013)
]
BooleanMask: TypeAlias = Union[
Sequence[bool],
_1DArray,
# IntoSeriesT, Would need support like (#2013)
]


SingleColSelector: TypeAlias = "SingleIndexSelector | SingleNameSelector"
MultiColSelector: TypeAlias = "MultiIndexSelector | MultiNameSelector | BooleanMask"

# NOTE: Not keeping, but using as a reference for each backend
Overload1: TypeAlias = "tuple[SingleIndexSelector, SingleColSelector]"
Overload2: TypeAlias = "str | tuple[MultiIndexSelector, SingleColSelector]"
Overload3: TypeAlias = """
SingleIndexSelector
| MultiIndexSelector
| MultiColSelector
| tuple[SingleIndexSelector, MultiColSelector]
| tuple[MultiIndexSelector, MultiColSelector]"""

OverloadUnion: TypeAlias = """
SingleIndexSelector
| SingleColSelector
| MultiColSelector
| MultiIndexSelector
| tuple[SingleIndexSelector, SingleColSelector]
| tuple[SingleIndexSelector, MultiColSelector]
| tuple[MultiIndexSelector, SingleColSelector]
| tuple[MultiIndexSelector, MultiColSelector]"""


class DTypes:
Decimal: type[dtypes.Decimal]
Int128: type[dtypes.Int128]
Expand Down

0 comments on commit 8d8895b

Please sign in to comment.