Merge pull request #17644 from BvB93/shape-base

ENH: Add annotations for `np.core.shape_base`

5 files changed, 103 insertions, 7 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index c8ef60c2d..612e56084 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -199,6 +199,16 @@ from numpy.core.numerictypes import (
     find_common_type as find_common_type,
 )
 
+from numpy.core.shape_base import (
+    atleast_1d as atleast_1d,
+    atleast_2d as atleast_2d,
+    atleast_3d as atleast_3d,
+    block as block,
+    hstack as hstack,
+    stack as stack,
+    vstack as vstack,
+)
+
 # Add an object to `__all__` if their stubs are defined in an external file;
 # their stubs will not be recognized otherwise.
 # NOTE: This is redundant for objects defined within this file.
@@ -262,15 +272,11 @@ asarray_chkfinite: Any
 asfarray: Any
 asmatrix: Any
 asscalar: Any
-atleast_1d: Any
-atleast_2d: Any
-atleast_3d: Any
 average: Any
 bartlett: Any
 bincount: Any
 bitwise_not: Any
 blackman: Any
-block: Any
 bmat: Any
 bool8: Any
 broadcast: Any
@@ -360,7 +366,6 @@ histogram2d: Any
 histogram_bin_edges: Any
 histogramdd: Any
 hsplit: Any
-hstack: Any
 i0: Any
 iinfo: Any
 imag: Any
@@ -494,7 +499,6 @@ singlecomplex: Any
 sort_complex: Any
 source: Any
 split: Any
-stack: Any
 string_: Any
 take_along_axis: Any
 tile: Any
@@ -527,7 +531,6 @@ vdot: Any
 vectorize: Any
 void0: Any
 vsplit: Any
-vstack: Any
 where: Any
 who: Any
 
diff --git a/numpy/core/shape_base.pyi b/numpy/core/shape_base.pyi
new file mode 100644
index 000000000..b20598b1a
--- /dev/null
+++ b/numpy/core/shape_base.pyi
@@ -0,0 +1,41 @@
+import sys
+from typing import TypeVar, overload, List, Sequence
+
+from numpy import ndarray
+from numpy.typing import ArrayLike
+
+if sys.version_info >= (3, 8):
+    from typing import SupportsIndex
+else:
+    from typing_extensions import Protocol
+    class SupportsIndex(Protocol):
+        def __index__(self) -> int: ...
+
+_ArrayType = TypeVar("_ArrayType", bound=ndarray)
+
+@overload
+def atleast_1d(__arys: ArrayLike) -> ndarray: ...
+@overload
+def atleast_1d(*arys: ArrayLike) -> List[ndarray]: ...
+
+@overload
+def atleast_2d(__arys: ArrayLike) -> ndarray: ...
+@overload
+def atleast_2d(*arys: ArrayLike) -> List[ndarray]: ...
+
+@overload
+def atleast_3d(__arys: ArrayLike) -> ndarray: ...
+@overload
+def atleast_3d(*arys: ArrayLike) -> List[ndarray]: ...
+
+def vstack(tup: Sequence[ArrayLike]) -> ndarray: ...
+def hstack(tup: Sequence[ArrayLike]) -> ndarray: ...
+@overload
+def stack(
+    arrays: Sequence[ArrayLike], axis: SupportsIndex = ..., out: None = ...
+) -> ndarray: ...
+@overload
+def stack(
+    arrays: Sequence[ArrayLike], axis: SupportsIndex = ..., out: _ArrayType = ...
+) -> _ArrayType: ...
+def block(arrays: ArrayLike) -> ndarray: ...
diff --git a/numpy/typing/tests/data/fail/array_constructors.py b/numpy/typing/tests/data/fail/array_constructors.py
index 5218572a6..9cb59fe5f 100644
--- a/numpy/typing/tests/data/fail/array_constructors.py
+++ b/numpy/typing/tests/data/fail/array_constructors.py
@@ -1,6 +1,7 @@
 import numpy as np
 
 a: np.ndarray
+generator = (i for i in range(10))
 
 np.require(a, requirements=1)  # E: No overload variant
 np.require(a, requirements="TEST")  # E: incompatible type
@@ -24,3 +25,7 @@ np.logspace(None, 'bob')  # E: Argument 1
 np.logspace(0, 2, base=None)  # E: Argument "base"
 
 np.geomspace(None, 'bob')  # E: Argument 1
+
+np.stack(generator)  # E: No overload variant
+np.hstack({1, 2})  # E: incompatible type
+np.vstack(1)  # E: incompatible type
diff --git a/numpy/typing/tests/data/pass/array_constructors.py b/numpy/typing/tests/data/pass/array_constructors.py
index 08f4b6211..63208f139 100644
--- a/numpy/typing/tests/data/pass/array_constructors.py
+++ b/numpy/typing/tests/data/pass/array_constructors.py
@@ -11,6 +11,7 @@ i8 = np.int64(1)
 
 A = np.array([1])
 B = A.view(SubClass).copy()
+B_stack = np.array([[1], [1]]).view(SubClass)
 C = [1]
 
 def func(i: int, j: int, **kwargs: Any) -> SubClass:
@@ -102,3 +103,26 @@ np.indices([1, 2, 3], sparse=True)
 np.fromfunction(func, (3, 5))
 
 np.identity(10)
+
+np.atleast_1d(C)
+np.atleast_1d(A)
+np.atleast_1d(C, C)
+np.atleast_1d(C, A)
+np.atleast_1d(A, A)
+
+np.atleast_2d(C)
+
+np.atleast_3d(C)
+
+np.vstack([C, C])
+np.vstack([C, A])
+np.vstack([A, A])
+
+np.hstack([C, C])
+
+np.stack([C, C])
+np.stack([C, C], axis=0)
+np.stack([C, C], out=B_stack)
+
+np.block([[C, C], [C, C]])
+np.block(A)
diff --git a/numpy/typing/tests/data/reveal/array_constructors.py b/numpy/typing/tests/data/reveal/array_constructors.py
index 04732ff2e..106174736 100644
--- a/numpy/typing/tests/data/reveal/array_constructors.py
+++ b/numpy/typing/tests/data/reveal/array_constructors.py
@@ -77,3 +77,26 @@ reveal_type(np.indices([1, 2, 3], sparse=True))  # E: tuple[numpy.ndarray]
 reveal_type(np.fromfunction(func, (3, 5)))  # E: SubClass
 
 reveal_type(np.identity(10))  # E: numpy.ndarray
+
+reveal_type(np.atleast_1d(A))  # E: numpy.ndarray
+reveal_type(np.atleast_1d(C))  # E: numpy.ndarray
+reveal_type(np.atleast_1d(A, A))  # E: list[numpy.ndarray]
+reveal_type(np.atleast_1d(A, C))  # E: list[numpy.ndarray]
+reveal_type(np.atleast_1d(C, C))  # E: list[numpy.ndarray]
+
+reveal_type(np.atleast_2d(A))  # E: numpy.ndarray
+
+reveal_type(np.atleast_3d(A))  # E: numpy.ndarray
+
+reveal_type(np.vstack([A, A]))  # E: numpy.ndarray
+reveal_type(np.vstack([A, C]))  # E: numpy.ndarray
+reveal_type(np.vstack([C, C]))  # E: numpy.ndarray
+
+reveal_type(np.hstack([A, A]))  # E: numpy.ndarray
+
+reveal_type(np.stack([A, A]))  # E: numpy.ndarray
+reveal_type(np.stack([A, A], axis=0))  # E: numpy.ndarray
+reveal_type(np.stack([A, A], out=B))  # E: SubClass
+
+reveal_type(np.block([[A, A], [A, A]]))  # E: numpy.ndarray
+reveal_type(np.block(C))  # E: numpy.ndarray