/
_array_like.py
156 lines (140 loc) · 4.06 KB
/
_array_like.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
from __future__ import annotations
# NOTE: Import `Sequence` from `typing` as we it is needed for a type-alias,
# not an annotation
from collections.abc import Collection, Callable
from typing import Any, Sequence, Protocol, Union, TypeVar
from numpy import (
ndarray,
dtype,
generic,
bool_,
unsignedinteger,
integer,
floating,
complexfloating,
number,
timedelta64,
datetime64,
object_,
void,
str_,
bytes_,
)
from ._nested_sequence import _NestedSequence
_T = TypeVar("_T")
_ScalarType = TypeVar("_ScalarType", bound=generic)
_DType = TypeVar("_DType", bound="dtype[Any]")
_DType_co = TypeVar("_DType_co", covariant=True, bound="dtype[Any]")
# The `_SupportsArray` protocol only cares about the default dtype
# (i.e. `dtype=None` or no `dtype` parameter at all) of the to-be returned
# array.
# Concrete implementations of the protocol are responsible for adding
# any and all remaining overloads
class _SupportsArray(Protocol[_DType_co]):
def __array__(self) -> ndarray[Any, _DType_co]: ...
class _SupportsArrayFunc(Protocol):
"""A protocol class representing `~class.__array_function__`."""
def __array_function__(
self,
func: Callable[..., Any],
types: Collection[type[Any]],
args: tuple[Any, ...],
kwargs: dict[str, Any],
) -> object: ...
# TODO: Wait until mypy supports recursive objects in combination with typevars
_FiniteNestedSequence = Union[
_T,
Sequence[_T],
Sequence[Sequence[_T]],
Sequence[Sequence[Sequence[_T]]],
Sequence[Sequence[Sequence[Sequence[_T]]]],
]
# A subset of `npt.ArrayLike` that can be parametrized w.r.t. `np.generic`
_ArrayLike = Union[
_SupportsArray["dtype[_ScalarType]"],
_NestedSequence[_SupportsArray["dtype[_ScalarType]"]],
]
# A union representing array-like objects; consists of two typevars:
# One representing types that can be parametrized w.r.t. `np.dtype`
# and another one for the rest
_DualArrayLike = Union[
_SupportsArray[_DType],
_NestedSequence[_SupportsArray[_DType]],
_T,
_NestedSequence[_T],
]
# TODO: support buffer protocols once
#
# https://bugs.python.org/issue27501
#
# is resolved. See also the mypy issue:
#
# https://github.com/python/typing/issues/593
ArrayLike = _DualArrayLike[
dtype,
Union[bool, int, float, complex, str, bytes],
]
# `ArrayLike<X>_co`: array-like objects that can be coerced into `X`
# given the casting rules `same_kind`
_ArrayLikeBool_co = _DualArrayLike[
"dtype[bool_]",
bool,
]
_ArrayLikeUInt_co = _DualArrayLike[
"dtype[Union[bool_, unsignedinteger[Any]]]",
bool,
]
_ArrayLikeInt_co = _DualArrayLike[
"dtype[Union[bool_, integer[Any]]]",
Union[bool, int],
]
_ArrayLikeFloat_co = _DualArrayLike[
"dtype[Union[bool_, integer[Any], floating[Any]]]",
Union[bool, int, float],
]
_ArrayLikeComplex_co = _DualArrayLike[
"dtype[Union[bool_, integer[Any], floating[Any], complexfloating[Any, Any]]]",
Union[bool, int, float, complex],
]
_ArrayLikeNumber_co = _DualArrayLike[
"dtype[Union[bool_, number[Any]]]",
Union[bool, int, float, complex],
]
_ArrayLikeTD64_co = _DualArrayLike[
"dtype[Union[bool_, integer[Any], timedelta64]]",
Union[bool, int],
]
_ArrayLikeDT64_co = Union[
_SupportsArray["dtype[datetime64]"],
_NestedSequence[_SupportsArray["dtype[datetime64]"]],
]
_ArrayLikeObject_co = Union[
_SupportsArray["dtype[object_]"],
_NestedSequence[_SupportsArray["dtype[object_]"]],
]
_ArrayLikeVoid_co = Union[
_SupportsArray["dtype[void]"],
_NestedSequence[_SupportsArray["dtype[void]"]],
]
_ArrayLikeStr_co = _DualArrayLike[
"dtype[str_]",
str,
]
_ArrayLikeBytes_co = _DualArrayLike[
"dtype[bytes_]",
bytes,
]
_ArrayLikeInt = _DualArrayLike[
"dtype[integer[Any]]",
int,
]
# Extra ArrayLike type so that pyright can deal with NDArray[Any]
# Used as the first overload, should only match NDArray[Any],
# not any actual types.
# https://github.com/numpy/numpy/pull/22193
class _UnknownType:
...
_ArrayLikeUnknown = _DualArrayLike[
"dtype[_UnknownType]",
_UnknownType,
]