Skip to content

Commit 689d05a

Browse files
jorenhamjorenham
committed
TYP: fix string-like ndarray rich comparison operators
1 parent d5fd966 commit 689d05a

File tree

2 files changed

+55
-24
lines changed

2 files changed

+55
-24
lines changed

‎numpy/__init__.pyi‎

Lines changed: 40 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@ from numpy._typing import (
2222
NDArray,
2323
_SupportsArray,
2424
_NestedSequence,
25-
_FiniteNestedSequence,
2625
_ArrayLike,
2726
_ArrayLikeBool_co,
2827
_ArrayLikeUInt_co,
@@ -33,28 +32,27 @@ from numpy._typing import (
3332
_ArrayLikeComplex128_co,
3433
_ArrayLikeComplex_co,
3534
_ArrayLikeNumber_co,
35+
_ArrayLikeObject_co,
36+
_ArrayLikeBytes_co,
37+
_ArrayLikeStr_co,
38+
_ArrayLikeString_co,
3639
_ArrayLikeTD64_co,
3740
_ArrayLikeDT64_co,
38-
_ArrayLikeObject_co,
39-
4041
# DTypes
4142
DTypeLike,
4243
_DTypeLike,
4344
_DTypeLikeVoid,
4445
_VoidDTypeLike,
45-
4646
# Shapes
4747
_Shape,
4848
_ShapeLike,
49-
5049
# Scalars
5150
_CharLike_co,
5251
_IntLike_co,
5352
_FloatLike_co,
5453
_TD64Like_co,
5554
_NumberLike_co,
5655
_ScalarLike_co,
57-
5856
# `number` precision
5957
NBitBase,
6058
# NOTE: Do not remove the extended precision bit-types even if seemingly unused;
@@ -77,7 +75,6 @@ from numpy._typing import (
7775
_NBitSingle,
7876
_NBitDouble,
7977
_NBitLongDouble,
80-
8178
# Character codes
8279
_BoolCodes,
8380
_UInt8Codes,
@@ -119,7 +116,6 @@ from numpy._typing import (
119116
_VoidCodes,
120117
_ObjectCodes,
121118
_StringCodes,
122-
123119
_UnsignedIntegerCodes,
124120
_SignedIntegerCodes,
125121
_IntegerCodes,
@@ -130,7 +126,6 @@ from numpy._typing import (
130126
_CharacterCodes,
131127
_FlexibleCodes,
132128
_GenericCodes,
133-
134129
# Ufuncs
135130
_UFunc_Nin1_Nout1,
136131
_UFunc_Nin2_Nout1,
@@ -2547,55 +2542,77 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
25472542
@overload # ?-d
25482543
def __iter__(self, /) -> Iterator[Any]: ...
25492544

2550-
# The last overload is for catching recursive objects whose
2551-
# nesting is too deep.
2552-
# The first overload is for catching `bytes` (as they are a subtype of
2553-
# `Sequence[int]`) and `str`. As `str` is a recursive sequence of
2554-
# strings, it will pass through the final overload otherwise
2555-
2545+
#
25562546
@overload
25572547
def __lt__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co, /) -> NDArray[np.bool]: ...
25582548
@overload
25592549
def __lt__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co, /) -> NDArray[np.bool]: ...
25602550
@overload
25612551
def __lt__(self: NDArray[datetime64], other: _ArrayLikeDT64_co, /) -> NDArray[np.bool]: ...
25622552
@overload
2563-
def __lt__(self: NDArray[object_], other: Any, /) -> NDArray[np.bool]: ...
2553+
def __lt__(self: NDArray[bytes_], other: _ArrayLikeBytes_co, /) -> NDArray[np.bool]: ...
25642554
@overload
2565-
def __lt__(self: NDArray[Any], other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
2555+
def __lt__(
2556+
self: ndarray[Any, dtype[str_] | dtypes.StringDType], other: _ArrayLikeStr_co | _ArrayLikeString_co, /
2557+
) -> NDArray[np.bool]: ...
2558+
@overload
2559+
def __lt__(self: NDArray[object_], other: object, /) -> NDArray[np.bool]: ...
2560+
@overload
2561+
def __lt__(self, other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
25662562

2563+
#
25672564
@overload
25682565
def __le__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co, /) -> NDArray[np.bool]: ...
25692566
@overload
25702567
def __le__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co, /) -> NDArray[np.bool]: ...
25712568
@overload
25722569
def __le__(self: NDArray[datetime64], other: _ArrayLikeDT64_co, /) -> NDArray[np.bool]: ...
25732570
@overload
2574-
def __le__(self: NDArray[object_], other: Any, /) -> NDArray[np.bool]: ...
2571+
def __le__(self: NDArray[bytes_], other: _ArrayLikeBytes_co, /) -> NDArray[np.bool]: ...
25752572
@overload
2576-
def __le__(self: NDArray[Any], other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
2573+
def __le__(
2574+
self: ndarray[Any, dtype[str_] | dtypes.StringDType], other: _ArrayLikeStr_co | _ArrayLikeString_co, /
2575+
) -> NDArray[np.bool]: ...
2576+
@overload
2577+
def __le__(self: NDArray[object_], other: object, /) -> NDArray[np.bool]: ...
2578+
@overload
2579+
def __le__(self, other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
25772580

2581+
#
25782582
@overload
25792583
def __gt__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co, /) -> NDArray[np.bool]: ...
25802584
@overload
25812585
def __gt__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co, /) -> NDArray[np.bool]: ...
25822586
@overload
25832587
def __gt__(self: NDArray[datetime64], other: _ArrayLikeDT64_co, /) -> NDArray[np.bool]: ...
25842588
@overload
2585-
def __gt__(self: NDArray[object_], other: Any, /) -> NDArray[np.bool]: ...
2589+
def __gt__(self: NDArray[bytes_], other: _ArrayLikeBytes_co, /) -> NDArray[np.bool]: ...
25862590
@overload
2587-
def __gt__(self: NDArray[Any], other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
2591+
def __gt__(
2592+
self: ndarray[Any, dtype[str_] | dtypes.StringDType], other: _ArrayLikeStr_co | _ArrayLikeString_co, /
2593+
) -> NDArray[np.bool]: ...
2594+
@overload
2595+
def __gt__(self: NDArray[object_], other: object, /) -> NDArray[np.bool]: ...
2596+
@overload
2597+
def __gt__(self, other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
25882598

2599+
#
25892600
@overload
25902601
def __ge__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co, /) -> NDArray[np.bool]: ...
25912602
@overload
25922603
def __ge__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co, /) -> NDArray[np.bool]: ...
25932604
@overload
25942605
def __ge__(self: NDArray[datetime64], other: _ArrayLikeDT64_co, /) -> NDArray[np.bool]: ...
25952606
@overload
2596-
def __ge__(self: NDArray[object_], other: Any, /) -> NDArray[np.bool]: ...
2607+
def __ge__(self: NDArray[bytes_], other: _ArrayLikeBytes_co, /) -> NDArray[np.bool]: ...
2608+
@overload
2609+
def __ge__(
2610+
self: ndarray[Any, dtype[str_] | dtypes.StringDType], other: _ArrayLikeStr_co | _ArrayLikeString_co, /
2611+
) -> NDArray[np.bool]: ...
2612+
@overload
2613+
def __ge__(self: NDArray[object_], other: object, /) -> NDArray[np.bool]: ...
25972614
@overload
2598-
def __ge__(self: NDArray[Any], other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
2615+
def __ge__(self, other: _ArrayLikeObject_co, /) -> NDArray[np.bool]: ...
25992616

26002617
# Unary ops
26012618

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from __future__ import annotations
22

3-
from typing import Any
3+
from typing import cast, Any
44
import numpy as np
55

66
c16 = np.complex128()
@@ -30,6 +30,9 @@
3030
AR_i: np.ndarray[Any, np.dtype[np.int_]] = np.array([1])
3131
AR_f: np.ndarray[Any, np.dtype[np.float64]] = np.array([1.0])
3232
AR_c: np.ndarray[Any, np.dtype[np.complex128]] = np.array([1.0j])
33+
AR_S: np.ndarray[Any, np.dtype[np.bytes_]] = np.array([b"a"], "S")
34+
AR_T = cast(np.ndarray[Any, np.dtypes.StringDType], np.array(["a"], "T"))
35+
AR_U: np.ndarray[Any, np.dtype[np.str_]] = np.array(["a"], "U")
3336
AR_m: np.ndarray[Any, np.dtype[np.timedelta64]] = np.array([np.timedelta64("1")])
3437
AR_M: np.ndarray[Any, np.dtype[np.datetime64]] = np.array([np.datetime64("1")])
3538
AR_O: np.ndarray[Any, np.dtype[np.object_]] = np.array([1], dtype=object)
@@ -66,6 +69,17 @@
6669
AR_c > AR_f
6770
AR_c > AR_c
6871

72+
AR_S > AR_S
73+
AR_S > b""
74+
75+
AR_T > AR_T
76+
AR_T > AR_U
77+
AR_T > ""
78+
79+
AR_U > AR_U
80+
AR_U > AR_T
81+
AR_U > ""
82+
6983
AR_m > AR_b
7084
AR_m > AR_u
7185
AR_m > AR_i

0 commit comments

Comments
 (0)