Skip to content

Commit 8c55907

Browse files
jorenhamjorenham
authored
TYP: ndarray.cum{prod,sum} shape-typing and improved dtypes (#31164)
1 parent cb0b583 commit 8c55907

4 files changed

Lines changed: 381 additions & 31 deletions

File tree

‎numpy/__init__.pyi‎

Lines changed: 182 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1799,7 +1799,7 @@ class _ArrayOrScalarCommon:
17991799
self, /, condition: _ArrayLikeInt_co, axis: SupportsIndex | None = None, *, out: ArrayT
18001800
) -> ArrayT: ...
18011801

1802-
# Keep in sync with `MaskedArray.cumprod`
1802+
#
18031803
@overload # out: None (default)
18041804
def cumprod(self, /, axis: SupportsIndex | None = None, dtype: DTypeLike | None = None, out: None = None) -> NDArray[Any]: ...
18051805
@overload # out: ndarray
@@ -1809,7 +1809,7 @@ class _ArrayOrScalarCommon:
18091809
self, /, axis: SupportsIndex | None = None, dtype: DTypeLike | None = None, *, out: ArrayT
18101810
) -> ArrayT: ...
18111811

1812-
# Keep in sync with `MaskedArray.cumsum`
1812+
#
18131813
@overload # out: None (default)
18141814
def cumsum(self, /, axis: SupportsIndex | None = None, dtype: DTypeLike | None = None, out: None = None) -> NDArray[Any]: ...
18151815
@overload # out: ndarray
@@ -2700,6 +2700,186 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeT_co, _DTypeT_co]):
27002700
where: _ArrayLikeBool_co | _NoValueType = ...,
27012701
) -> ndarray[_ShapeT_co]: ...
27022702

2703+
# keep in sync with `MaskedArray.cumprod`
2704+
@override # type: ignore[override]
2705+
@overload # number | object_
2706+
def cumprod[DTypeT: dtype[number | object_]](
2707+
self: ndarray[Any, DTypeT],
2708+
axis: None = None,
2709+
dtype: None = None,
2710+
out: None = None,
2711+
) -> ndarray[_1D, DTypeT]: ...
2712+
@overload # bool_
2713+
def cumprod(
2714+
self: NDArray[bool_],
2715+
axis: None = None,
2716+
dtype: None = None,
2717+
out: None = None,
2718+
) -> ndarray[_1D, dtype[int_]]: ...
2719+
@overload # dtype: <known> (keyword)
2720+
def cumprod[ScalarT: generic](
2721+
self: NDArray[number | bool_ | object_],
2722+
axis: None = None,
2723+
*,
2724+
dtype: _DTypeLike[ScalarT],
2725+
out: None = None,
2726+
) -> ndarray[_1D, dtype[ScalarT]]: ...
2727+
@overload # dtype: <unknown> (keyword)
2728+
def cumprod(
2729+
self: NDArray[number | bool_ | object_],
2730+
axis: None = None,
2731+
*,
2732+
dtype: DTypeLike,
2733+
out: None = None,
2734+
) -> ndarray[_1D]: ...
2735+
@overload # dtype: <known> (positional)
2736+
def cumprod[ScalarT: generic](
2737+
self: NDArray[number | bool_ | object_],
2738+
axis: None,
2739+
dtype: _DTypeLike[ScalarT],
2740+
out: None = None,
2741+
) -> ndarray[_1D, dtype[ScalarT]]: ...
2742+
@overload # dtype: <unknown> (positional)
2743+
def cumprod(
2744+
self: NDArray[number | bool_ | object_],
2745+
axis: None,
2746+
dtype: DTypeLike,
2747+
out: None = None,
2748+
) -> ndarray[_1D]: ...
2749+
@overload # axis: <given>
2750+
def cumprod[ArrayT: NDArray[number | object_]](
2751+
self: ArrayT,
2752+
axis: SupportsIndex,
2753+
dtype: None = None,
2754+
out: None = None,
2755+
) -> ArrayT: ...
2756+
@overload # bool_, axis: <given>
2757+
def cumprod[ShapeT: _Shape](
2758+
self: ndarray[ShapeT, dtype[bool_]],
2759+
axis: SupportsIndex,
2760+
dtype: None = None,
2761+
out: None = None,
2762+
) -> ndarray[ShapeT, dtype[int_]]: ...
2763+
@overload # axis: <given>, dtype: <known>
2764+
def cumprod[ShapeT: _Shape, ScalarT: generic](
2765+
self: ndarray[ShapeT, dtype[number | bool_ | object_]],
2766+
axis: SupportsIndex,
2767+
dtype: _DTypeLike[ScalarT],
2768+
out: None = None,
2769+
) -> ndarray[ShapeT, dtype[ScalarT]]: ...
2770+
@overload # axis: <given>, dtype: <unknown>
2771+
def cumprod[ShapeT: _Shape](
2772+
self: ndarray[ShapeT, dtype[number | bool_ | object_]],
2773+
axis: SupportsIndex,
2774+
dtype: DTypeLike,
2775+
out: None = None,
2776+
) -> ndarray[ShapeT]: ...
2777+
@overload # out: ndarray
2778+
def cumprod[ArrayT: ndarray](
2779+
self: NDArray[number | bool_ | object_],
2780+
axis: SupportsIndex | None,
2781+
dtype: DTypeLike | None,
2782+
out: ArrayT,
2783+
) -> ArrayT: ...
2784+
@overload
2785+
def cumprod[ArrayT: ndarray]( # pyright: ignore[reportIncompatibleMethodOverride]
2786+
self: NDArray[number | bool_ | object_],
2787+
axis: SupportsIndex | None = None,
2788+
dtype: DTypeLike | None = None,
2789+
*,
2790+
out: ArrayT,
2791+
) -> ArrayT: ...
2792+
2793+
# keep in sync with `MaskedArray.cumsum`
2794+
@override # type: ignore[override]
2795+
@overload # number | timedelta64 | object_
2796+
def cumsum[DTypeT: dtype[number | timedelta64 | object_]](
2797+
self: ndarray[Any, DTypeT],
2798+
axis: None = None,
2799+
dtype: None = None,
2800+
out: None = None,
2801+
) -> ndarray[_1D, DTypeT]: ...
2802+
@overload # bool_
2803+
def cumsum(
2804+
self: NDArray[bool_],
2805+
axis: None = None,
2806+
dtype: None = None,
2807+
out: None = None,
2808+
) -> ndarray[_1D, dtype[int_]]: ...
2809+
@overload # dtype: <known> (keyword)
2810+
def cumsum[ScalarT: generic](
2811+
self: NDArray[number | bool_ | timedelta64 | object_],
2812+
axis: None = None,
2813+
*,
2814+
dtype: _DTypeLike[ScalarT],
2815+
out: None = None,
2816+
) -> ndarray[_1D, dtype[ScalarT]]: ...
2817+
@overload # dtype: <unknown> (keyword)
2818+
def cumsum(
2819+
self: NDArray[number | bool_ | timedelta64 | object_],
2820+
axis: None = None,
2821+
*,
2822+
dtype: DTypeLike,
2823+
out: None = None,
2824+
) -> ndarray[_1D]: ...
2825+
@overload # dtype: <known> (positional)
2826+
def cumsum[ScalarT: generic](
2827+
self: NDArray[number | bool_ | timedelta64 | object_],
2828+
axis: None,
2829+
dtype: _DTypeLike[ScalarT],
2830+
out: None = None,
2831+
) -> ndarray[_1D, dtype[ScalarT]]: ...
2832+
@overload # dtype: <unknown> (positional)
2833+
def cumsum(
2834+
self: NDArray[number | bool_ | timedelta64 | object_],
2835+
axis: None,
2836+
dtype: DTypeLike,
2837+
out: None = None,
2838+
) -> ndarray[_1D]: ...
2839+
@overload # axis: <given>
2840+
def cumsum[ArrayT: NDArray[number | timedelta64 | object_]](
2841+
self: ArrayT,
2842+
axis: SupportsIndex,
2843+
dtype: None = None,
2844+
out: None = None,
2845+
) -> ArrayT: ...
2846+
@overload # bool_, axis: <given>
2847+
def cumsum[ShapeT: _Shape](
2848+
self: ndarray[ShapeT, dtype[bool_]],
2849+
axis: SupportsIndex,
2850+
dtype: None = None,
2851+
out: None = None,
2852+
) -> ndarray[ShapeT, dtype[int_]]: ...
2853+
@overload # axis: <given>, dtype: <known>
2854+
def cumsum[ShapeT: _Shape, ScalarT: generic](
2855+
self: ndarray[ShapeT, dtype[number | bool_ | timedelta64 | object_]],
2856+
axis: SupportsIndex,
2857+
dtype: _DTypeLike[ScalarT],
2858+
out: None = None,
2859+
) -> ndarray[ShapeT, dtype[ScalarT]]: ...
2860+
@overload # axis: <given>, dtype: <unknown>
2861+
def cumsum[ShapeT: _Shape](
2862+
self: ndarray[ShapeT, dtype[number | bool_ | timedelta64 | object_]],
2863+
axis: SupportsIndex,
2864+
dtype: DTypeLike,
2865+
out: None = None,
2866+
) -> ndarray[ShapeT]: ...
2867+
@overload # out: ndarray
2868+
def cumsum[ArrayT: ndarray](
2869+
self: NDArray[number | bool_ | timedelta64 | object_],
2870+
axis: SupportsIndex | None,
2871+
dtype: DTypeLike | None,
2872+
out: ArrayT,
2873+
) -> ArrayT: ...
2874+
@overload
2875+
def cumsum[ArrayT: ndarray]( # pyright: ignore[reportIncompatibleMethodOverride]
2876+
self: NDArray[number | bool_ | timedelta64 | object_],
2877+
axis: SupportsIndex | None = None,
2878+
dtype: DTypeLike | None = None,
2879+
*,
2880+
out: ArrayT,
2881+
) -> ArrayT: ...
2882+
27032883
#
27042884
@override # type: ignore[override]
27052885
@overload # +integer | ~object_

0 commit comments

Comments
 (0)