Skip to content

Commit 4964ca8

Browse files
jorenhamcharris
jorenham
authored and
charris
committed
TYP: isclose shape-typing fix for 2d array-likes (#32205)
1 parent c37ed94 commit 4964ca8

2 files changed

Lines changed: 20 additions & 16 deletions

File tree

‎numpy/_core/numeric.pyi‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1280,58 +1280,58 @@ def allclose(
12801280
def isclose(
12811281
a: _NumberLike_co,
12821282
b: _NumberLike_co,
1283-
rtol: ArrayLike = 1e-5,
1284-
atol: ArrayLike = 1e-8,
1283+
rtol: _ScalarLike_co = 1e-5,
1284+
atol: _ScalarLike_co = 1e-8,
12851285
equal_nan: py_bool = False,
12861286
) -> np.bool: ...
12871287
@overload # known shape, same shape or scalar
12881288
def isclose[ShapeT: _Shape](
12891289
a: np.ndarray[ShapeT],
12901290
b: np.ndarray[ShapeT] | _NumberLike_co,
1291-
rtol: ArrayLike = 1e-5,
1292-
atol: ArrayLike = 1e-8,
1291+
rtol: _ScalarLike_co = 1e-5,
1292+
atol: _ScalarLike_co = 1e-8,
12931293
equal_nan: py_bool = False,
12941294
) -> np.ndarray[ShapeT, np.dtype[np.bool]]: ...
12951295
@overload # same shape or scalar, known shape
12961296
def isclose[ShapeT: _Shape](
12971297
a: np.ndarray[ShapeT] | _NumberLike_co,
12981298
b: np.ndarray[ShapeT],
1299-
rtol: ArrayLike = 1e-5,
1300-
atol: ArrayLike = 1e-8,
1299+
rtol: _ScalarLike_co = 1e-5,
1300+
atol: _ScalarLike_co = 1e-8,
13011301
equal_nan: py_bool = False,
13021302
) -> np.ndarray[ShapeT, np.dtype[np.bool]]: ...
13031303
@overload # 1d sequence, <=1d array-like
13041304
def isclose(
13051305
a: Sequence[_NumberLike_co],
13061306
b: Sequence[_NumberLike_co] | _NumberLike_co | np.ndarray[tuple[int]],
1307-
rtol: ArrayLike = 1e-5,
1308-
atol: ArrayLike = 1e-8,
1307+
rtol: _ScalarLike_co = 1e-5,
1308+
atol: _ScalarLike_co = 1e-8,
13091309
equal_nan: py_bool = False,
13101310
) -> np.ndarray[tuple[int], np.dtype[np.bool]]: ...
13111311
@overload # <=1d array-like, 1d sequence
13121312
def isclose(
13131313
a: Sequence[_NumberLike_co] | _NumberLike_co | np.ndarray[tuple[int]],
13141314
b: Sequence[_NumberLike_co],
1315-
rtol: ArrayLike = 1e-5,
1316-
atol: ArrayLike = 1e-8,
1315+
rtol: _ScalarLike_co = 1e-5,
1316+
atol: _ScalarLike_co = 1e-8,
13171317
equal_nan: py_bool = False,
13181318
) -> np.ndarray[tuple[int], np.dtype[np.bool]]: ...
13191319
@overload # 2d sequence, <=2d array-like
13201320
def isclose(
13211321
a: Sequence[Sequence[_NumberLike_co]],
13221322
b: Sequence[Sequence[_NumberLike_co]] | Sequence[_NumberLike_co] | _NumberLike_co | np.ndarray[tuple[int] | tuple[int, int]],
1323-
rtol: ArrayLike = 1e-5,
1324-
atol: ArrayLike = 1e-8,
1323+
rtol: _ScalarLike_co = 1e-5,
1324+
atol: _ScalarLike_co = 1e-8,
13251325
equal_nan: py_bool = False,
1326-
) -> np.ndarray[tuple[int], np.dtype[np.bool]]: ...
1326+
) -> np.ndarray[tuple[int, int], np.dtype[np.bool]]: ...
13271327
@overload # <=2d array-like, 2d sequence
13281328
def isclose(
13291329
b: Sequence[Sequence[_NumberLike_co]] | Sequence[_NumberLike_co] | _NumberLike_co | np.ndarray[tuple[int] | tuple[int, int]],
13301330
a: Sequence[Sequence[_NumberLike_co]],
1331-
rtol: ArrayLike = 1e-5,
1332-
atol: ArrayLike = 1e-8,
1331+
rtol: _ScalarLike_co = 1e-5,
1332+
atol: _ScalarLike_co = 1e-8,
13331333
equal_nan: py_bool = False,
1334-
) -> np.ndarray[tuple[int], np.dtype[np.bool]]: ...
1334+
) -> np.ndarray[tuple[int, int], np.dtype[np.bool]]: ...
13351335
@overload # unknown shape, unknown shape
13361336
def isclose(
13371337
a: ArrayLike,

‎numpy/typing/tests/data/reveal/numeric.pyi‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ _sub_nd_i8: SubClass
2727
_to_1d_bool: list[bool]
2828
_to_1d_int: list[int]
2929
_to_1d_float: list[float]
30+
_to_2d_float: list[list[float]]
3031
_to_1d_complex: list[complex]
3132

3233
###
@@ -161,6 +162,9 @@ assert_type(np.allclose(AR_i8, AR_i8), bool)
161162
assert_type(np.isclose(i8, i8), np.bool)
162163
assert_type(np.isclose(i8, AR_i8), npt.NDArray[np.bool])
163164
assert_type(np.isclose(_to_1d_int, _to_1d_int), np.ndarray[tuple[int], np.dtype[np.bool]])
165+
assert_type(np.isclose(_to_1d_int, _to_2d_float), np.ndarray[tuple[int, int], np.dtype[np.bool]])
166+
assert_type(np.isclose(_to_2d_float, _to_1d_int), np.ndarray[tuple[int, int], np.dtype[np.bool]])
167+
assert_type(np.isclose(_to_2d_float, _to_2d_float), np.ndarray[tuple[int, int], np.dtype[np.bool]])
164168
assert_type(np.isclose(AR_i8, AR_i8), npt.NDArray[np.bool])
165169

166170
assert_type(np.array_equal(i8, AR_i8), bool)

0 commit comments

Comments
 (0)