Merge pull request numpy#27157 from guan404ming/mean-type-overload

TYP: add td64 overload for `mean`

Author: jorenham

numpy/_core/fromnumeric.pyi

@@ -11,6 +11,7 @@ from numpy import (
     float16,
     floating,
     complexfloating,
+    timedelta64,
     object_,
     generic,
     _OrderKACF,
@@ -35,6 +36,7 @@ from numpy._typing import (
     _ArrayLikeFloat_co,
     _ArrayLikeComplex_co,
     _ArrayLikeObject_co,
+    _ArrayLikeTD64_co,
     _IntLike_co,
     _BoolLike_co,
     _ComplexLike_co,
@@ -1062,6 +1064,16 @@ def mean(
     where: _ArrayLikeBool_co = ...,
 ) -> complexfloating[Any, Any]: ...
 @overload
+def mean(
+    a: _ArrayLikeTD64_co,
+    axis: None = ...,
+    dtype: None = ...,
+    out: None = ...,
+    keepdims: Literal[False] = ...,
+    *,
+    where: _ArrayLikeBool_co = ...,
+) -> timedelta64: ...
+@overload
 def mean(
     a: _ArrayLikeComplex_co | _ArrayLikeObject_co,
     axis: None | _ShapeLike = ...,

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

@@ -6,6 +6,7 @@ import numpy.typing as npt
 A = np.array(True, ndmin=2, dtype=bool)
 A.setflags(write=False)
 AR_U: npt.NDArray[np.str_]
+AR_M: npt.NDArray[np.datetime64]
 
 a = np.bool(True)
 
@@ -147,6 +148,7 @@ np.mean(a, axis=1.0)  # E: No overload variant
 np.mean(a, out=False)  # E: No overload variant
 np.mean(a, keepdims=1.0)  # E: No overload variant
 np.mean(AR_U)  # E: incompatible type
+np.mean(AR_M)  # E: incompatible type
 
 np.std(a, axis=1.0)  # E: No overload variant
 np.std(a, out=False)  # E: No overload variant

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

@@ -21,6 +21,7 @@ AR_u8: npt.NDArray[np.uint64]
 AR_i8: npt.NDArray[np.int64]
 AR_O: npt.NDArray[np.object_]
 AR_subclass: NDArraySubclass
+AR_m: npt.NDArray[np.timedelta64]
 
 b: np.bool
 f4: np.float32
@@ -294,6 +295,7 @@ assert_type(np.around(AR_f4, out=AR_subclass), NDArraySubclass)
 assert_type(np.mean(AR_b), np.floating[Any])
 assert_type(np.mean(AR_i8), np.floating[Any])
 assert_type(np.mean(AR_f4), np.floating[Any])
+assert_type(np.mean(AR_m), np.timedelta64)
 assert_type(np.mean(AR_c16), np.complexfloating[Any, Any])
 assert_type(np.mean(AR_O), Any)
 assert_type(np.mean(AR_f4, axis=0), Any)