8000 Merge pull request #18377 from BvB93/ufunclike · numpy/numpy@f6a71fe · GitHub
[go: up one dir, main page]
Skip to content

Commit f6a71fe

Browse files
charrischarris
authored
Merge pull request #18377 from BvB93/ufunclike
ENH: Add annotations for `np.lib.ufunclike`
2 parents 3823775 + 692f00b commit f6a71fe

File tree

7 files changed

+164
-8
lines changed

7 files changed

+164
-8
lines changed

‎numpy/__init__.pyi

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,12 @@ from numpy.core.shape_base import (
332332
vstack as vstack,
333333
)
334334

335+
from numpy.lib.ufunclike import (
336+
fix as fix,
337+
isposinf as isposinf,
338+
isneginf as isneginf,
339+
)
340+
335341
__all__: List[str]
336342
__path__: List[str]
337343
__version__: str
@@ -401,7 +407,6 @@ extract: Any
401407
eye: Any
402408
fill_diagonal: Any
403409
finfo: Any
404-
fix: Any
405410
flip: Any
406411
fliplr: Any
407412
flipud: Any
@@ -437,8 +442,6 @@ is_busday: Any
437442
iscomplex: Any
438443
iscomplexobj: Any
439444
isin: Any
440-
isneginf: Any
441-
isposinf: Any
442445
isreal: Any
443446
isrealobj: Any
444447
iterable: Any

‎numpy/lib/__init__.pyi

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,11 @@
11
from typing import Any, List
22

3+
from numpy.lib.ufunclike import (
4+
fix as fix,
5+
isposinf as isposinf,
6+
isneginf as isneginf,
7+
)
8+
39
__all__: List[str]
410

511
emath: Any
@@ -108,9 +114,6 @@ tril_indices: Any
108114
tril_indices_from: Any
109115
triu_indices: Any
110116
triu_indices_from: Any
111-
fix: Any
112-
isneginf: Any
113-
isposinf: Any
114117
pad: Any
115118
poly: Any
116119
roots: Any

‎numpy/lib/ufunclike.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,8 @@ def isposinf(x, out=None):
189189
try:
190190
signbit = ~nx.signbit(x)
191191
except TypeError as e:
192-
raise TypeError('This operation is not supported for complex values '
192+
dtype = nx.asanyarray(x).dtype
193+
raise TypeError(f'This operation is not supported for {dtype} values '
193194
'because it would be ambiguous.') from e
194195
else:
195196
return nx.logical_and(is_inf, signbit, out)
@@ -260,7 +261,8 @@ def isneginf(x, out=None):
260261
try:
261262
signbit = nx.signbit(x)
262263
except TypeError as e:
263-
raise TypeError('This operation is not supported for complex values '
264+
dtype = nx.asanyarray(x).dtype
265+
raise TypeError(f'This operation is not supported for {dtype} values '
264266
'because it would be ambiguous.') from e
265267
else:
266268
return nx.logical_and(is_inf, signbit, out)

‎numpy/lib/ufunclike.pyi

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
from typing import Any, overload, TypeVar, List, Union
2+
3+
from numpy import floating, bool_, ndarray
4+
from numpy.typing import (
5+
_ArrayLikeFloat_co,
6+
_ArrayLikeObject_co,
7+
_ArrayOrScalar,
8+
)
9+
10+
_ArrayType = TypeVar("_ArrayType", bound=ndarray[Any, Any])
11+
12+
__all__: List[str]
13+
14+
@overload
15+
def fix(
16+
x: _ArrayLikeFloat_co,
17+
out: None = ...,
18+
) -> _ArrayOrScalar[floating[Any]]: ...
19+
@overload
20+
def fix(
21+
x: _ArrayLikeObject_co,
22+
out: None = ...,
23+
) -> Any: ...
24+
@overload
25+
def fix(
26+
x: Union[_ArrayLikeFloat_co, _ArrayLikeObject_co],
27+
out: _ArrayType,
28+
) -> _ArrayType: ...
29+
30+
@overload
31+
def isposinf(
32+
x: _ArrayLikeFloat_co,
33+
out: None = ...,
34+
) -> _ArrayOrScalar[bool_]: ...
35+
@overload
36+
def isposinf(
37+
x: _ArrayLikeFloat_co,
38+
out: _ArrayType,
39+
) -> _ArrayType: ...
40+
41+
@overload
42+
def isneginf(
43+
x: _ArrayLikeFloat_co,
44+
out: None = ...,
45+
) -> _ArrayOrScalar[bool_]: ...
46+
@overload
47+
def isneginf(
48+
x: _ArrayLikeFloat_co,
49+
out: _ArrayType,
50+
) -> _ArrayType: ...

‎numpy/typing/tests/data/fail/ufunclike.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
from typing import List, Any
2+
import numpy as np
3+
4+
AR_c: np.ndarray[Any, np.dtype[np.complex128]]
5+
AR_m: np.ndarray[Any, np.dtype[np.timedelta64]]
6+
AR_M: np.ndarray[Any, np.dtype[np.datetime64]]
7+
AR_O: np.ndarray[Any, np.dtype[np.object_]]
8+
9+
np.fix(AR_c) # E: incompatible type
10+
np.fix(AR_m) # E: incompatible type
11+
np.fix(AR_M) # E: incompatible type
12+
13+
np.isposinf(AR_c) # E: incompatible type
14+
np.isposinf(AR_m) # E: incompatible type
15+
np.isposinf(AR_M) # E: incompatible type
16+
np.isposinf(AR_O) # E: incompatible type
17+
18+
np.isneginf(AR_c) # E: incompatible type
19+
np.isneginf(AR_m) # E: incompatible type
20+
np.isneginf(AR_M) # E: incompatible type
21+
np.isneginf(AR_O) # E: incompatible type

‎numpy/typing/tests/data/pass/ufunclike.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
from __future__ import annotations
2+
from typing import Any
3+
import numpy as np
4+
5+
6+
class Object:
7+
def __ceil__(self) -> Object:
8+
return self
9+
10+
def __floor__(self) -> Object:
11+
return self
12+
13+
def __ge__(self, value: object) -> bool:
14+
return True
15+
16+
def __array__(
17+
self, dtype: None = None
18+
) -> np.ndarray[Any, np.dtype[np.object_]]:
19+
ret = np.empty((), dtype=object)
20+
ret[()] = self
21+
return ret
22+
23+
24+
AR_LIKE_b = [True, True, False]
25+
AR_LIKE_u = [np.uint32(1), np.uint32(2), np.uint32(3)]
26+
AR_LIKE_i = [1, 2, 3]
27+
AR_LIKE_f = [1.0, 2.0, 3.0]
28+
AR_LIKE_O = [Object(), Object(), Object()]
29+
AR_U: np.ndarray[Any, np.dtype[np.str_]] = np.zeros(3, dtype="U5")
30+
31+
np.fix(AR_LIKE_b)
32+
np.fix(AR_LIKE_u)
33+
np.fix(AR_LIKE_i)
34+
np.fix(AR_LIKE_f)
35+
np.fix(AR_LIKE_O)
36+
np.fix(AR_LIKE_f, out=AR_U)
37+
38+
np.isposinf(AR_LIKE_b)
39+
np.isposinf(AR_LIKE_u)
40+
np.isposinf(AR_LIKE_i)
41+
np.isposinf(AR_LIKE_f)
42+
np.isposinf(AR_LIKE_f, out=AR_U)
43+
44+
np.isneginf(AR_LIKE_b)
45+
np.isneginf(AR_LIKE_u)
46+
np.isneginf(AR_LIKE_i)
47+
np.isneginf(AR_LIKE_f)
48+
np.isneginf(AR_LIKE_f, out=AR_U)

‎numpy/typing/tests/data/reveal/ufunclike.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from typing import List, Any
2+
import numpy as np
3+
4+
AR_LIKE_b: List[bool]
5+
AR_LIKE_u: List[np.uint32]
6+
AR_LIKE_i: List[int]
7+
AR_LIKE_f: List[float]
8+
AR_LIKE_O: List[np.object_]
9+
10+
AR_U: np.ndarray[Any, np.dtype[np.str_]]
11+
12+
reveal_type(np.fix(AR_LIKE_b)) # E: Union[numpy.floating[Any], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]]
13+
reveal_type(np.fix(AR_LIKE_u)) # E: Union[numpy.floating[Any], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]]
14+
reveal_type(np.fix(AR_LIKE_i)) # E: Union[numpy.floating[Any], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]]
15+
reveal_type(np.fix(AR_LIKE_f)) # E: Union[numpy.floating[Any], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]]
16+
reveal_type(np.fix(AR_LIKE_O)) # E: Any
17+
reveal_type(np.fix(AR_LIKE_f, out=AR_U)) # E: numpy.ndarray[Any, numpy.dtype[numpy.str_]]
18+
19+
reveal_type(np.isposinf(AR_LIKE_b)) # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]]
20+
reveal_type(np.isposinf(AR_LIKE_u)) # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]]
21+
reveal_type(np.isposinf(AR_LIKE_i)) # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]]
22+
reveal_type(np.isposinf(AR_LIKE_f)) # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]]
23+
reveal_type(np.isposinf(AR_LIKE_f, out=AR_U)) # E: numpy.ndarray[Any, numpy.dtype[numpy.str_]]
24+
25+
reveal_type(np.isneginf(AR_LIKE_b)) # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]]
26+
reveal_type(np.isneginf(AR_LIKE_u)) # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]]
27+
reveal_type(np.isneginf(AR_LIKE_i)) # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]]
28+
reveal_type(np.isneginf(AR_LIKE_f)) # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]]
29+
reveal_type(np.isneginf(AR_LIKE_f, out=AR_U)) # E: numpy.ndarray[Any, numpy.dtype[numpy.str_]]

0 commit comments

Comments
 (0)
0