Skip to content

Commit b25ecaa

Browse files
charrischarris
authored
Merge pull request numpy#27756 from jorenham/typing/1d-shape-typing
TYP: 1-d shape-typing for ``ndarray.flatten`` and ``ravel``
2 parents b79f912 + 547deac commit b25ecaa

File tree

6 files changed

+60
-53
lines changed

6 files changed

+60
-53
lines changed

‎numpy/__init__.pyi‎

Lines changed: 12 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -2193,17 +2193,8 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType_co, _DType_co]):
21932193
axis: None | SupportsIndex = ...,
21942194
) -> ndarray[_Shape, _DType_co]: ...
21952195

2196-
# TODO: use `tuple[int]` as shape type once covariant (#26081)
2197-
def flatten(
2198-
self,
2199-
order: _OrderKACF = ...,
2200-
) -> ndarray[_Shape, _DType_co]: ...
2201-
2202-
# TODO: use `tuple[int]` as shape type once covariant (#26081)
2203-
def ravel(
2204-
self,
2205-
order: _OrderKACF = ...,
2206-
) -> ndarray[_Shape, _DType_co]: ...
2196+
def flatten(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], _DType_co]: ...
2197+
def ravel(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], _DType_co]: ...
22072198

22082199
@overload
22092200
def reshape(
@@ -3099,11 +3090,10 @@ _NBit_fc = TypeVar("_NBit_fc", _NBitHalf, _NBitSingle, _NBitDouble, _NBitLongDou
30993090
class generic(_ArrayOrScalarCommon):
31003091
@abstractmethod
31013092
def __init__(self, *args: Any, **kwargs: Any) -> None: ...
3102-
# TODO: use `tuple[()]` as shape type once covariant (#26081)
31033093
@overload
3104-
def __array__(self, dtype: None = ..., /) -> NDArray[Self]: ...
3094+
def __array__(self, dtype: None = None, /) -> ndarray[tuple[()], dtype[Self]]: ...
31053095
@overload
3106-
def __array__(self, dtype: _DType, /) -> ndarray[_Shape, _DType]: ...
3096+
def __array__(self, dtype: _DType, /) -> ndarray[tuple[()], _DType]: ...
31073097
def __hash__(self) -> int: ...
31083098
@property
31093099
def base(self) -> None: ...
@@ -3117,7 +3107,7 @@ class generic(_ArrayOrScalarCommon):
31173107
def strides(self) -> tuple[()]: ...
31183108
def byteswap(self, inplace: L[False] = ...) -> Self: ...
31193109
@property
3120-
def flat(self) -> flatiter[NDArray[Self]]: ...
3110+
def flat(self) -> flatiter[ndarray[tuple[int], dtype[Self]]]: ...
31213111

31223112
if sys.version_info >= (3, 12):
31233113
def __buffer__(self, flags: int, /) -> memoryview: ...
@@ -3201,8 +3191,8 @@ class generic(_ArrayOrScalarCommon):
32013191
) -> _NdArraySubClass: ...
32023192

32033193
def repeat(self, repeats: _ArrayLikeInt_co, axis: None | SupportsIndex = ...) -> NDArray[Self]: ...
3204-
def flatten(self, order: _OrderKACF = ...) -> NDArray[Self]: ...
3205-
def ravel(self, order: _OrderKACF = ...) -> NDArray[Self]: ...
3194+
def flatten(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], dtype[Self]]: ...
3195+
def ravel(self, /, order: _OrderKACF = "C") -> ndarray[tuple[int], dtype[Self]]: ...
32063196

32073197
@overload
32083198
def reshape(self, shape: _ShapeLike, /, *, order: _OrderACF = ...) -> NDArray[Self]: ...
@@ -4480,13 +4470,12 @@ class poly1d:
44804470
@coefficients.setter
44814471
def coefficients(self, value: NDArray[Any]) -> None: ...
44824472

4483-
__hash__: ClassVar[None] # type: ignore
4473+
__hash__: ClassVar[None] # type: ignore[assignment] # pyright: ignore[reportIncompatibleMethodOverride]
44844474

4485-
# TODO: use `tuple[int]` as shape type once covariant (#26081)
44864475
@overload
4487-
def __array__(self, t: None = ..., copy: None | bool = ...) -> NDArray[Any]: ...
4476+
def __array__(self, /, t: None = None, copy: builtins.bool | None = None) -> ndarray[tuple[int], dtype[Any]]: ...
44884477
@overload
4489-
def __array__(self, t: _DType, copy: None | bool = ...) -> ndarray[_Shape, _DType]: ...
4478+
def __array__(self, /, t: _DType, copy: builtins.bool | None = None) -> ndarray[tuple[int], _DType]: ...
44904479

44914480
@overload
44924481
def __call__(self, val: _ScalarLike_co) -> Any: ...
@@ -4656,8 +4645,8 @@ class matrix(ndarray[_Shape2DType_co, _DType_co]):
46564645

46574646
def squeeze(self, axis: None | _ShapeLike = ...) -> matrix[_Shape2D, _DType_co]: ...
46584647
def tolist(self: _SupportsItem[_T]) -> list[list[_T]]: ...
4659-
def ravel(self, order: _OrderKACF = ...) -> matrix[_Shape2D, _DType_co]: ...
4660-
def flatten(self, order: _OrderKACF = ...) -> matrix[_Shape2D, _DType_co]: ...
4648+
def ravel(self, /, order: _OrderKACF = "C") -> matrix[tuple[L[1], int], _DType_co]: ... # pyright: ignore[reportIncompatibleMethodOverride]
4649+
def flatten(self, /, order: _OrderKACF = "C") -> matrix[tuple[L[1], int], _DType_co]: ... # pyright: ignore[reportIncompatibleMethodOverride]
46614650

46624651
@property
46634652
def T(self) -> matrix[_Shape2D, _DType_co]: ...

‎numpy/_core/fromnumeric.pyi‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ from numpy._typing import (
3939
ArrayLike,
4040
_ArrayLike,
4141
NDArray,
42+
_NestedSequence,
4243
_ShapeLike,
4344
_ArrayLikeBool_co,
4445
_ArrayLikeUInt_co,
@@ -438,10 +439,27 @@ def trace(
438439
out: _ArrayType = ...,
439440
) -> _ArrayType: ...
440441

442+
_Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_SCT]]
443+
444+
@overload
445+
def ravel(a: _ArrayLike[_SCT], order: _OrderKACF = "C") -> _Array1D[_SCT]: ...
446+
@overload
447+
def ravel(a: bytes | _NestedSequence[bytes], order: _OrderKACF = "C") -> _Array1D[np.bytes_]: ...
448+
@overload
449+
def ravel(a: str | _NestedSequence[str], order: _OrderKACF = "C") -> _Array1D[np.str_]: ...
450+
@overload
451+
def ravel(a: bool | _NestedSequence[bool], order: _OrderKACF = "C") -> _Array1D[np.bool]: ...
452+
@overload
453+
def ravel(a: int | _NestedSequence[int], order: _OrderKACF = "C") -> _Array1D[np.int_ | np.bool]: ...
454+
@overload
455+
def ravel(a: float | _NestedSequence[float], order: _OrderKACF = "C") -> _Array1D[np.float64 | np.int_ | np.bool]: ...
441456
@overload
442-
def ravel(a: _ArrayLike[_SCT], order: _OrderKACF = ...) -> NDArray[_SCT]: ...
457+
def ravel(
458+
a: complex | _NestedSequence[complex],
459+
order: _OrderKACF = "C",
460+
) -> _Array1D[np.complex128 | np.float64 | np.int_ | np.bool]: ...
443461
@overload
444-
def ravel(a: ArrayLike, order: _OrderKACF = ...) -> NDArray[Any]: ...
462+
def ravel(a: ArrayLike, order: _OrderKACF = "C") -> np.ndarray[tuple[int], np.dtype[Any]]: ...
445463

446464
@overload
447465
def nonzero(a: np.generic | np.ndarray[tuple[()], Any]) -> NoReturn: ...

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -121,11 +121,11 @@ assert_type(np.trace(AR_b), Any)
121121
assert_type(np.trace(AR_f4), Any)
122122
assert_type(np.trace(AR_f4, out=AR_subclass), NDArraySubclass)
123123

124-
assert_type(np.ravel(b), npt.NDArray[np.bool])
125-
assert_type(np.ravel(f4), npt.NDArray[np.float32])
126-
assert_type(np.ravel(f), npt.NDArray[Any])
127-
assert_type(np.ravel(AR_b), npt.NDArray[np.bool])
128-
assert_type(np.ravel(AR_f4), npt.NDArray[np.float32])
124+
assert_type(np.ravel(b), np.ndarray[tuple[int], np.dtype[np.bool]])
125+
assert_type(np.ravel(f4), np.ndarray[tuple[int], np.dtype[np.float32]])
126+
assert_type(np.ravel(f), np.ndarray[tuple[int], np.dtype[np.float64 | np.int_ | np.bool]])
127+
assert_type(np.ravel(AR_b), np.ndarray[tuple[int], np.dtype[np.bool]])
128+
assert_type(np.ravel(AR_f4), np.ndarray[tuple[int], np.dtype[np.float32]])
129129

130130
assert_type(np.nonzero(b), NoReturn)
131131
assert_type(np.nonzero(f4), NoReturn)

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -173,11 +173,11 @@ assert_type(AR_f8.trace(out=B), SubClass)
173173
assert_type(AR_f8.item(), float)
174174
assert_type(AR_U.item(), str)
175175

176-
assert_type(AR_f8.ravel(), npt.NDArray[np.float64])
177-
assert_type(AR_U.ravel(), npt.NDArray[np.str_])
176+
assert_type(AR_f8.ravel(), np.ndarray[tuple[int], np.dtype[np.float64]])
177+
assert_type(AR_U.ravel(), np.ndarray[tuple[int], np.dtype[np.str_]])
178178

179-
assert_type(AR_f8.flatten(), npt.NDArray[np.float64])
180-
assert_type(AR_U.flatten(), npt.NDArray[np.str_])
179+
assert_type(AR_f8.flatten(), np.ndarray[tuple[int], np.dtype[np.float64]])
180+
assert_type(AR_U.flatten(), np.ndarray[tuple[int], np.dtype[np.str_]])
181181

182182
assert_type(AR_f8.reshape(1), npt.NDArray[np.float64])
183183
assert_type(AR_U.reshape(1), npt.NDArray[np.str_])

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,12 @@ assert_type(nd.transpose((1, 0)), npt.NDArray[np.int64])
2525
assert_type(nd.swapaxes(0, 1), npt.NDArray[np.int64])
2626

2727
# flatten
28-
assert_type(nd.flatten(), npt.NDArray[np.int64])
29-
assert_type(nd.flatten("C"), npt.NDArray[np.int64])
28+
assert_type(nd.flatten(), np.ndarray[tuple[int], np.dtype[np.int64]])
29+
assert_type(nd.flatten("C"), np.ndarray[tuple[int], np.dtype[np.int64]])
3030

3131
# ravel
32-
assert_type(nd.ravel(), npt.NDArray[np.int64])
33-
assert_type(nd.ravel("C"), npt.NDArray[np.int64])
32+
assert_type(nd.ravel(), np.ndarray[tuple[int], np.dtype[np.int64]])
33+
assert_type(nd.ravel("C"), np.ndarray[tuple[int], np.dtype[np.int64]])
3434

3535
# squeeze
3636
assert_type(nd.squeeze(), npt.NDArray[np.int64])

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

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -93,21 +93,21 @@ assert_type(c16.tolist(), complex)
9393
assert_type(U.tolist(), str)
9494
assert_type(S.tolist(), bytes)
9595

96-
assert_type(b.ravel(), npt.NDArray[np.bool])
97-
assert_type(i8.ravel(), npt.NDArray[np.int64])
98-
assert_type(u8.ravel(), npt.NDArray[np.uint64])
99-
assert_type(f8.ravel(), npt.NDArray[np.float64])
100-
assert_type(c16.ravel(), npt.NDArray[np.complex128])
101-
assert_type(U.ravel(), npt.NDArray[np.str_])
102-
assert_type(S.ravel(), npt.NDArray[np.bytes_])
103-
104-
assert_type(b.flatten(), npt.NDArray[np.bool])
105-
assert_type(i8.flatten(), npt.NDArray[np.int64])
106-
assert_type(u8.flatten(), npt.NDArray[np.uint64])
107-
assert_type(f8.flatten(), npt.NDArray[np.float64])
108-
assert_type(c16.flatten(), npt.NDArray[np.complex128])
109-
assert_type(U.flatten(), npt.NDArray[np.str_])
110-
assert_type(S.flatten(), npt.NDArray[np.bytes_])
96+
assert_type(b.ravel(), np.ndarray[tuple[int], np.dtype[np.bool]])
97+
assert_type(i8.ravel(), np.ndarray[tuple[int], np.dtype[np.int64]])
98+
assert_type(u8.ravel(), np.ndarray[tuple[int], np.dtype[np.uint64]])
99+
assert_type(f8.ravel(), np.ndarray[tuple[int], np.dtype[np.float64]])
100+
assert_type(c16.ravel(), np.ndarray[tuple[int], np.dtype[np.complex128]])
101+
assert_type(U.ravel(), np.ndarray[tuple[int], np.dtype[np.str_]])
102+
assert_type(S.ravel(), np.ndarray[tuple[int], np.dtype[np.bytes_]])
103+
104+
assert_type(b.flatten(), np.ndarray[tuple[int], np.dtype[np.bool]])
105+
assert_type(i8.flatten(), np.ndarray[tuple[int], np.dtype[np.int64]])
106+
assert_type(u8.flatten(), np.ndarray[tuple[int], np.dtype[np.uint64]])
107+
assert_type(f8.flatten(), np.ndarray[tuple[int], np.dtype[np.float64]])
108+
assert_type(c16.flatten(), np.ndarray[tuple[int], np.dtype[np.complex128]])
109+
assert_type(U.flatten(), np.ndarray[tuple[int], np.dtype[np.str_]])
110+
assert_type(S.flatten(), np.ndarray[tuple[int], np.dtype[np.bytes_]])
111111

112112
assert_type(b.reshape(1), npt.NDArray[np.bool])
113113
assert_type(i8.reshape(1), npt.NDArray[np.int64])

0 commit comments

Comments
 (0)