|
1 | 1 | from __future__ import annotations |
2 | | -from typing import TypeAlias, TYPE_CHECKING |
| 2 | +from typing import TypeAlias, TYPE_CHECKING, Protocol, Any, TypeVar, Generic |
3 | 3 | from datetime import date, datetime, time, timedelta |
4 | 4 | from decimal import Decimal |
5 | 5 | from uuid import UUID |
6 | | -from collections.abc import Mapping |
| 6 | +from collections.abc import Mapping, Iterator |
7 | 7 |
|
8 | 8 | if TYPE_CHECKING: |
9 | 9 | from ._expression import Expression |
10 | 10 |
|
| 11 | +_T_co = TypeVar("_T_co", covariant=True) |
| 12 | +_S_co = TypeVar("_S_co", bound=tuple[Any, ...], covariant=True) |
| 13 | +_D_co = TypeVar("_D_co", covariant=True) |
| 14 | + |
| 15 | +class NPTypeLike(Protocol, Generic[_T_co]): ... |
| 16 | + |
| 17 | +class NPArrayLike(Protocol, Generic[_S_co, _D_co]): |
| 18 | + def __len__(self) -> int: ... |
| 19 | + def __contains__(self, value: object, /) -> bool: ... |
| 20 | + def __iter__(self) -> Iterator[_D_co]: ... |
| 21 | + def __array__(self, *args: Any, **kwargs: Any) -> Any: ... |
| 22 | + def __array_finalize__(self, *args: Any, **kwargs: Any) -> None: ... |
| 23 | + def __array_wrap__(self, *args: Any, **kwargs: Any) -> Any: ... |
| 24 | + def __getitem__(self, *args: Any, **kwargs: Any) -> Any: ... |
| 25 | + def __setitem__(self, *args: Any, **kwargs: Any) -> None: ... |
| 26 | + @property |
| 27 | + def shape(self) -> _S_co: ... |
| 28 | + @property |
| 29 | + def dtype(self) -> Any: ... |
| 30 | + @property |
| 31 | + def ndim(self) -> int: ... |
| 32 | + @property |
| 33 | + def size(self) -> int: ... |
| 34 | + |
11 | 35 | NumericLiteral: TypeAlias = int | float | Decimal |
12 | 36 | """Python objects that can be converted to a numerical `ConstantExpression` (integer or floating points numbers.)""" |
13 | 37 | TemporalLiteral: TypeAlias = date | datetime | time | timedelta |
|
19 | 43 | """ |
20 | 44 | NonNestedLiteral: TypeAlias = NumericLiteral | TemporalLiteral | str | bool | BlobLiteral | UUID |
21 | 45 | PythonLiteral: TypeAlias = ( |
22 | | - NonNestedLiteral | list[PythonLiteral] | tuple[PythonLiteral, ...] | dict[PythonLiteral, PythonLiteral] | None |
| 46 | + NonNestedLiteral |
| 47 | + | list[PythonLiteral] |
| 48 | + | tuple[PythonLiteral, ...] |
| 49 | + | dict[PythonLiteral, PythonLiteral] |
| 50 | + | NPArrayLike[Any, Any] |
| 51 | + | None |
23 | 52 | ) |
24 | 53 | """Python objects that can be converted to a `ConstantExpression`.""" |
25 | 54 | # the field_ids argument to to_parquet and write_parquet has a recursive structure |
|
0 commit comments