8000 TYP: preserve shape-type in ndarray.astype() · numpy/numpy@b866d68 · GitHub
[go: up one dir, main page]
Skip to content

Commit b866d68

Browse files
ntrrgcntrrgc
committed
TYP: preserve shape-type in ndarray.astype()
This patch changes the return type in astype() from NDArray to ndarray so that shape information is preserved and adds tests for it. Similar changes are added to np.astype() for consistency.
1 parent bbf4836 commit b866d68

File tree

3 files changed

+14
-6
lines changed

3 files changed

+14
-6
lines changed

‎numpy/__init__.pyi

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2496,7 +2496,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DType_co]):
24962496
casting: _CastingKind = ...,
24972497
subok: builtins.bool = ...,
24982498
copy: builtins.bool | _CopyMode = ...,
2499-
) -> NDArray[_SCT]: ...
2499+
) -> ndarray[_ShapeT_co, dtype[_SCT]]: ...
25002500
@overload
25012501
def astype(
25022502
self,
@@ -2505,7 +2505,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DType_co]):
25052505
casting: _CastingKind = ...,
25062506
subok: builtins.bool = ...,
25072507
copy: builtins.bool | _CopyMode = ...,
2508-
) -> NDArray[Any]: ...
2508+
) -> ndarray[_ShapeT_co, dtype[Any]]: ...
25092509

25102510
@overload
25112511
def view(self) -> Self: ...

‎numpy/_core/numeric.pyi

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -872,15 +872,15 @@ def array_equiv(a1: ArrayLike, a2: ArrayLike) -> bool: ...
872872

873873
@overload
874874
def astype(
875-
x: NDArray[Any],
875+
x: ndarray[_ShapeType, dtype[Any]],
876876
dtype: _DTypeLike[_SCT],
877877
copy: bool = ...,
878878
device: None | L["cpu"] = ...,
879-
) -> NDArray[_SCT]: ...
879+
) -> ndarray[_ShapeType, dtype[_SCT]]: ...
880880
@overload
881881
def astype(
882-
x: NDArray[Any],
882+
x: ndarray[_ShapeType, dtype[Any]],
883883
dtype: DTypeLike,
884884
copy: bool = ...,
885885
device: None | L["cpu"] = ...,
886-
) -> NDArray[Any]: ...
886+
) -> ndarray[_ShapeType, dtype[Any]]: ...

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ i4_2d: np.ndarray[tuple[int, int], np.dtype[np.int32]]
1111
f8_3d: np.ndarray[tuple[int, int, int], np.dtype[np.float64]]
1212
cG_4d: np.ndarray[tuple[int, int, int, int], np.dtype[np.clongdouble]]
1313
i0_nd: npt.NDArray[np.int_]
14+
uncertain_dtype: np.int32 | np.float64 | np.str_
1415

1516
# item
1617
assert_type(i0_nd.item(), int)
@@ -50,6 +51,13 @@ assert_type(i0_nd.astype(np.float64, "K", "unsafe", True, True), npt.NDArray[np.
5051

5152
assert_type(np.astype(i0_nd, np.float64), npt.NDArray[np.float64])
5253

54+
assert_type(i4_2d.astype(np.uint16), np.ndarray[tuple[int, int], np.dtype[np.uint16]])
55+
assert_type(np.astype(i4_2d, np.uint16), np.ndarray[tuple[int, int], np.dtype[np.uint16]])
56+
assert_type(f8_3d.astype(np.int16), np.ndarray[tuple[int, int, int], np.dtype[np.int16]])
57+
assert_type(np.astype(f8_3d, np.int16), np.ndarray[tuple[int, int, int], np.dtype[np.int16]])
58+
assert_type(i4_2d.astype(uncertain_dtype), np.ndarray[tuple[int, int], np.dtype[np.generic[Any]]])
59+
assert_type(np.astype(i4_2d, uncertain_dtype), np.ndarray[tuple[int, int], np.dtype[Any]])
60+
5361
# byteswap
5462
assert_type(i0_nd.byteswap(), npt.NDArray[np.int_])
5563
assert_type(i0_nd.byteswap(True), npt.NDArray[np.int_])

0 commit comments

Comments
 (0)
0