Merge pull request numpy#27683 from bersbersbers/bersbersbers-27638

TYP: Improve `np.sum` and `np.mean` return types with given `dtype`

Author: jorenham

numpy/_core/fromnumeric.pyi

@@ -586,6 +586,16 @@ def clip(
     casting: _CastingKind = ...,
 ) -> _ArrayType: ...
 
+@overload
+def sum(
+    a: _ArrayLike[_SCT],
+    axis: None = ...,
+    dtype: None = ...,
+    out: None  = ...,
+    keepdims: Literal[False] = ...,
+    initial: _NumberLike_co = ...,
+    where: _ArrayLikeBool_co = ...,
+) -> _SCT: ...
 @overload
 def sum(
     a: _ArrayLike[_SCT],
@@ -595,8 +605,50 @@ def sum(
     keepdims: bool = ...,
     initial: _NumberLike_co = ...,
     where: _ArrayLikeBool_co = ...,
+) -> _SCT | NDArray[_SCT]: ...
+@overload
+def sum(
+    a: ArrayLike,
+    axis: None,
+    dtype: _DTypeLike[_SCT],
+    out: None = ...,
+    keepdims: Literal[False] = ...,
+    initial: _NumberLike_co = ...,
+    where: _ArrayLikeBool_co = ...,
+) -> _SCT: ...
+@overload
+def sum(
+    a: ArrayLike,
+    axis: None = ...,
+    *,
+    dtype: _DTypeLike[_SCT],
+    out: None = ...,
+    keepdims: Literal[False] = ...,
+    initial: _NumberLike_co = ...,
+    where: _ArrayLikeBool_co = ...,
 ) -> _SCT: ...
 @overload
+def sum(
+    a: ArrayLike,
+    axis: None | _ShapeLike,
+    dtype: _DTypeLike[_SCT],
+    out: None  = ...,
+    keepdims: bool = ...,
+    initial: _NumberLike_co = ...,
+    where: _ArrayLikeBool_co = ...,
+) -> _SCT | NDArray[_SCT]: ...
+@overload
+def sum(
+    a: ArrayLike,
+    axis: None | _ShapeLike = ...,
+    *,
+    dtype: _DTypeLike[_SCT],
+    out: None  = ...,
+    keepdims: bool = ...,
+    initial: _NumberLike_co = ...,
+    where: _ArrayLikeBool_co = ...,
+) -> _SCT | NDArray[_SCT]: ...
+@overload
 def sum(
     a: ArrayLike,
     axis: None | _ShapeLike = ...,
@@ -1207,6 +1259,26 @@ def mean(
     where: _ArrayLikeBool_co = ...,
 ) -> _SCT: ...
 @overload
+def mean(
+    a: _ArrayLikeComplex_co | _ArrayLikeObject_co,
+    axis: None,
+    dtype: _DTypeLike[_SCT],
+    out: None = ...,
+    keepdims: bool = ...,
+    *,
+    where: _ArrayLikeBool_co = ...,
+) -> _SCT | NDArray[_SCT]: ...
+@overload
+def mean(
+    a: _ArrayLikeComplex_co | _ArrayLikeObject_co,
+    axis: None = ...,
+    *,
+    dtype: _DTypeLike[_SCT],
+    out: None = ...,
+    keepdims: bool = ...,
+    where: _ArrayLikeBool_co = ...,
+) -> _SCT | NDArray[_SCT]: ...
+@overload
 def mean(
     a: _ArrayLikeComplex_co | _ArrayLikeObject_co,
     axis: None | _ShapeLike = ...,

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

@@ -169,6 +169,12 @@ assert_type(np.sum(AR_f4), np.float32)
 assert_type(np.sum(AR_b, axis=0), Any)
 assert_type(np.sum(AR_f4, axis=0), Any)
 assert_type(np.sum(AR_f4, out=AR_subclass), NDArraySubclass)
+assert_type(np.sum(AR_f4, dtype=np.float64), np.float64)
+assert_type(np.sum(AR_f4, None, np.float64), np.float64)
+assert_type(np.sum(AR_f4, dtype=np.float64, keepdims=False), np.float64)
+assert_type(np.sum(AR_f4, None, np.float64, keepdims=False), np.float64)
+assert_type(np.sum(AR_f4, dtype=np.float64, keepdims=True), np.float64 | npt.NDArray[np.float64])
+assert_type(np.sum(AR_f4, None, np.float64, keepdims=True), np.float64 | npt.NDArray[np.float64])
 
 assert_type(np.all(b), np.bool)
 assert_type(np.all(f4), np.bool)
@@ -310,6 +316,12 @@ assert_type(np.mean(AR_f4, keepdims=True), Any)
 assert_type(np.mean(AR_f4, dtype=float), Any)
 assert_type(np.mean(AR_f4, dtype=np.float64), np.float64)
 assert_type(np.mean(AR_f4, out=AR_subclass), NDArraySubclass)
+assert_type(np.mean(AR_f4, dtype=np.float64), np.float64)
+assert_type(np.mean(AR_f4, None, np.float64), np.float64)
+assert_type(np.mean(AR_f4, dtype=np.float64, keepdims=False), np.float64)
+assert_type(np.mean(AR_f4, None, np.float64, keepdims=False), np.float64)
+assert_type(np.mean(AR_f4, dtype=np.float64, keepdims=True), np.float64 | npt.NDArray[np.float64])
+assert_type(np.mean(AR_f4, None, np.float64, keepdims=True), np.float64 | npt.NDArray[np.float64])
 
 assert_type(np.std(AR_b), np.floating[Any])
 assert_type(np.std(AR_i8), np.floating[Any])