Merge pull request #17644 from BvB93/shape-base

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

3 files changed, 52 insertions, 0 deletions

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