TYP,ENH: Improve the dtype-overload of stack, hstack and vstack

Xref numpy#21627
BvB93 · Nov 14, 2022 · 473b80f · 473b80f
1 parent 8205e5d
commit 473b80f
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 11 deletions.
diff --git a/numpy/core/shape_base.pyi b/numpy/core/shape_base.pyi
@@ -2,7 +2,13 @@ from collections.abc import Sequence
 from typing import TypeVar, overload, Any, SupportsIndex
 
 from numpy import generic, _CastingKind
-from numpy._typing import ArrayLike, NDArray, _ArrayLike, DTypeLike
+from numpy._typing import (
+    NDArray,
+    ArrayLike,
+    DTypeLike,
+    _ArrayLike,
+    _DTypeLike,
+)
 
 _SCT = TypeVar("_SCT", bound=generic)
 _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
@@ -34,39 +40,62 @@ def atleast_3d(*arys: ArrayLike) -> list[NDArray[Any]]: ...
 def vstack(
     tup: Sequence[_ArrayLike[_SCT]],
     *,
-    dtype: DTypeLike  = ..., 
+    dtype: None = ...,
     casting: _CastingKind = ...
 ) -> NDArray[_SCT]: ...
 @overload
 def vstack(
     tup: Sequence[ArrayLike],
     *,
-    dtype: DTypeLike  = ..., 
+    dtype: _DTypeLike[_SCT],
+    casting: _CastingKind = ...
+) -> NDArray[_SCT]: ...
+@overload
+def vstack(
+    tup: Sequence[ArrayLike],
+    *,
+    dtype: DTypeLike = ...,
     casting: _CastingKind = ...
 ) -> NDArray[Any]: ...
 
 @overload
 def hstack(
-    tup: Sequence[_ArrayLike[_SCT]], 
+    tup: Sequence[_ArrayLike[_SCT]],
+    *,
+    dtype: None = ...,
+    casting: _CastingKind = ...
+) -> NDArray[_SCT]: ...
+@overload
+def hstack(
+    tup: Sequence[ArrayLike],
     *,
-    dtype: DTypeLike  = ..., 
+    dtype: _DTypeLike[_SCT],
     casting: _CastingKind = ...
 ) -> NDArray[_SCT]: ...
 @overload
 def hstack(
-    tup: Sequence[ArrayLike], 
+    tup: Sequence[ArrayLike],
     *,
-    dtype: DTypeLike  = ..., 
+    dtype: DTypeLike = ...,
     casting: _CastingKind = ...
 ) -> NDArray[Any]: ...
 
 @overload
 def stack(
     arrays: Sequence[_ArrayLike[_SCT]],
     axis: SupportsIndex = ...,
-    out: None  = ...,
+    out: None = ...,
     *,
-    dtype: DTypeLike = ...,
+    dtype: None = ...,
+    casting: _CastingKind = ...
+) -> NDArray[_SCT]: ...
+@overload
+def stack(
+    arrays: Sequence[ArrayLike],
+    axis: SupportsIndex = ...,
+    out: None = ...,
+    *,
+    dtype: _DTypeLike[_SCT],
     casting: _CastingKind = ...
 ) -> NDArray[_SCT]: ...
 @overload
@@ -75,7 +104,7 @@ def stack(
     axis: SupportsIndex = ...,
     out: None = ...,
     *,
-    dtype: DTypeLike  = ...,
+    dtype: DTypeLike = ...,
     casting: _CastingKind = ...
 ) -> NDArray[Any]: ...
 @overload
@@ -84,7 +113,7 @@ def stack(
     axis: SupportsIndex = ...,
     out: _ArrayType = ...,
     *,
-    dtype: DTypeLike  = ...,
+    dtype: DTypeLike = ...,
     casting: _CastingKind = ...
 ) -> _ArrayType: ...
 

diff --git a/numpy/typing/tests/data/reveal/array_constructors.pyi b/numpy/typing/tests/data/reveal/array_constructors.pyi
@@ -187,12 +187,15 @@ reveal_type(np.atleast_2d(A))  # E: ndarray[Any, dtype[{float64}]]
 reveal_type(np.atleast_3d(A))  # E: ndarray[Any, dtype[{float64}]]
 
 reveal_type(np.vstack([A, A]))  # E: ndarray[Any, Any]
+reveal_type(np.vstack([A, A], dtype=np.float64))  # E: ndarray[Any, dtype[{float64}]]
 reveal_type(np.vstack([A, C]))  # E: ndarray[Any, dtype[Any]]
 reveal_type(np.vstack([C, C]))  # E: ndarray[Any, dtype[Any]]
 
 reveal_type(np.hstack([A, A]))  # E: ndarray[Any, Any]
+reveal_type(np.hstack([A, A], dtype=np.float64))  # E: ndarray[Any, dtype[{float64}]]
 
 reveal_type(np.stack([A, A]))  # E: Any
+reveal_type(np.stack([A, A], dtype=np.float64))  # E: ndarray[Any, dtype[{float64}]]
 reveal_type(np.stack([A, C]))  # E: ndarray[Any, dtype[Any]]
 reveal_type(np.stack([C, C]))  # E: ndarray[Any, dtype[Any]]
 reveal_type(np.stack([A, A], axis=0))  # E: Any