Merge pull request #17326 from BvB93/asarray

ENH: Add annotations for five array construction functions

10 files changed, 208 insertions, 54 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 6a1c5b72b..30d15ed12 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -107,6 +107,14 @@ from numpy.core.fromnumeric import (
     var,
 )
 
+from numpy.core._asarray import (
+    asarray as asarray,
+    asanyarray as asanyarray,
+    ascontiguousarray as ascontiguousarray,
+    asfortranarray as asfortranarray,
+    require as require,
+)
+
 # 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.
@@ -169,12 +177,8 @@ array2string: Any
 array_repr: Any
 array_split: Any
 array_str: Any
-asanyarray: Any
-asarray: Any
 asarray_chkfinite: Any
-ascontiguousarray: Any
 asfarray: Any
-asfortranarray: Any
 asmatrix: Any
 asscalar: Any
 atleast_1d: Any
@@ -396,7 +400,6 @@ recarray: Any
 recfromcsv: Any
 recfromtxt: Any
 record: Any
-require: Any
 result_type: Any
 roots: Any
 rot90: Any
diff --git a/numpy/core/_asarray.pyi b/numpy/core/_asarray.pyi
new file mode 100644
index 000000000..e074d69d2
--- /dev/null
+++ b/numpy/core/_asarray.pyi
@@ -0,0 +1,77 @@
+import sys
+from typing import TypeVar, Union, Iterable, overload
+
+from numpy import ndarray, _OrderKACF
+from numpy.typing import ArrayLike, DtypeLike
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+_ArrayType = TypeVar("_ArrayType", bound=ndarray)
+
+def asarray(
+    a: object,
+    dtype: DtypeLike = ...,
+    order: _OrderKACF = ...,
+    *,
+    like: ArrayLike = ...
+) -> ndarray: ...
+@overload
+def asanyarray(
+    a: _ArrayType,
+    dtype: None = ...,
+    order: _OrderKACF = ...,
+    *,
+    like: ArrayLike = ...
+) -> _ArrayType: ...
+@overload
+def asanyarray(
+    a: object,
+    dtype: DtypeLike = ...,
+    order: _OrderKACF = ...,
+    *,
+    like: ArrayLike = ...
+) -> ndarray: ...
+def ascontiguousarray(
+    a: object, dtype: DtypeLike = ..., *, like: ArrayLike = ...
+) -> ndarray: ...
+def asfortranarray(
+    a: object, dtype: DtypeLike = ..., *, like: ArrayLike = ...
+) -> ndarray: ...
+
+_Requirements = Literal[
+    "C", "C_CONTIGUOUS", "CONTIGUOUS",
+    "F", "F_CONTIGUOUS", "FORTRAN",
+    "A", "ALIGNED",
+    "W", "WRITEABLE",
+    "O", "OWNDATA"
+]
+_E = Literal["E", "ENSUREARRAY"]
+_RequirementsWithE = Union[_Requirements, _E]
+
+@overload
+def require(
+    a: _ArrayType,
+    dtype: None = ...,
+    requirements: Union[None, _Requirements, Iterable[_Requirements]] = ...,
+    *,
+    like: ArrayLike = ...
+) -> _ArrayType: ...
+@overload
+def require(
+    a: object,
+    dtype: DtypeLike = ...,
+    requirements: Union[_E, Iterable[_RequirementsWithE]] = ...,
+    *,
+    like: ArrayLike = ...
+) -> ndarray: ...
+@overload
+def require(
+    a: object,
+    dtype: DtypeLike = ...,
+    requirements: Union[None, _Requirements, Iterable[_Requirements]] = ...,
+    *,
+    like: ArrayLike = ...
+) -> ndarray: ...
diff --git a/numpy/typing/tests/data/fail/linspace.py b/numpy/typing/tests/data/fail/array_constructors.py
index a9769c5d6..5218572a6 100644
--- a/numpy/typing/tests/data/fail/linspace.py
+++ b/numpy/typing/tests/data/fail/array_constructors.py
@@ -1,5 +1,18 @@
 import numpy as np
 
+a: np.ndarray
+
+np.require(a, requirements=1)  # E: No overload variant
+np.require(a, requirements="TEST")  # E: incompatible type
+
+np.zeros("test")  # E: incompatible type
+np.zeros()  # E: Too few arguments
+
+np.ones("test")  # E: incompatible type
+np.ones()  # E: Too few arguments
+
+np.array(0, float, True)  # E: Too many positional
+
 np.linspace(None, 'bob')  # E: No overload variant
 np.linspace(0, 2, num=10.0)  # E: No overload variant
 np.linspace(0, 2, endpoint='True')  # E: No overload variant
diff --git a/numpy/typing/tests/data/fail/simple.py b/numpy/typing/tests/data/fail/simple.py
deleted file mode 100644
index 57c08fb7d..000000000
--- a/numpy/typing/tests/data/fail/simple.py
+++ /dev/null
@@ -1,12 +0,0 @@
-"""Simple expression that should fail with mypy."""
-
-import numpy as np
-
-# Array creation routines checks
-np.zeros("test")  # E: incompatible type
-np.zeros()  # E: Too few arguments
-
-np.ones("test")  # E: incompatible type
-np.ones()  # E: Too few arguments
-
-np.array(0, float, True)  # E: Too many positional
diff --git a/numpy/typing/tests/data/pass/array_constructors.py b/numpy/typing/tests/data/pass/array_constructors.py
new file mode 100644
index 000000000..bf29e52b9
--- /dev/null
+++ b/numpy/typing/tests/data/pass/array_constructors.py
@@ -0,0 +1,66 @@
+from typing import List
+import numpy as np
+
+class Index:
+    def __index__(self) -> int:
+        return 0
+
+class SubClass(np.ndarray): ...
+
+A = np.array([1])
+B = A.view(SubClass).copy()
+C = [1]
+
+np.array(1, dtype=float)
+np.array(1, copy=False)
+np.array(1, order='F')
+np.array(1, order=None)
+np.array(1, subok=True)
+np.array(1, ndmin=3)
+np.array(1, str, copy=True, order='C', subok=False, ndmin=2)
+
+np.asarray(A)
+np.asarray(B)
+np.asarray(C)
+
+np.asanyarray(A)
+np.asanyarray(B)
+np.asanyarray(B, dtype=int)
+np.asanyarray(C)
+
+np.ascontiguousarray(A)
+np.ascontiguousarray(B)
+np.ascontiguousarray(C)
+
+np.asfortranarray(A)
+np.asfortranarray(B)
+np.asfortranarray(C)
+
+np.require(A)
+np.require(B)
+np.require(B, dtype=int)
+np.require(B, requirements=None)
+np.require(B, requirements="E")
+np.require(B, requirements=["ENSUREARRAY"])
+np.require(B, requirements={"F", "E"})
+np.require(B, requirements=["C", "OWNDATA"])
+np.require(B, requirements="W")
+np.require(B, requirements="A")
+np.require(C)
+
+np.linspace(0, 2)
+np.linspace(0.5, [0, 1, 2])
+np.linspace([0, 1, 2], 3)
+np.linspace(0j, 2)
+np.linspace(0, 2, num=10)
+np.linspace(0, 2, endpoint=True)
+np.linspace(0, 2, retstep=True)
+np.linspace(0j, 2j, retstep=True)
+np.linspace(0, 2, dtype=bool)
+np.linspace([0, 1], [2, 3], axis=Index())
+
+np.logspace(0, 2, base=2)
+np.logspace(0, 2, base=2)
+np.logspace(0, 2, base=[1j, 2j], num=2)
+
+np.geomspace(1, 2)
diff --git a/numpy/typing/tests/data/pass/linspace.py b/numpy/typing/tests/data/pass/linspace.py
deleted file mode 100644
index 8c6d0d56b..000000000
--- a/numpy/typing/tests/data/pass/linspace.py
+++ /dev/null
@@ -1,22 +0,0 @@
-import numpy as np
-
-class Index:
-    def __index__(self) -> int:
-        return 0
-
-np.linspace(0, 2)
-np.linspace(0.5, [0, 1, 2])
-np.linspace([0, 1, 2], 3)
-np.linspace(0j, 2)
-np.linspace(0, 2, num=10)
-np.linspace(0, 2, endpoint=True)
-np.linspace(0, 2, retstep=True)
-np.linspace(0j, 2j, retstep=True)
-np.linspace(0, 2, dtype=bool)
-np.linspace([0, 1], [2, 3], axis=Index())
-
-np.logspace(0, 2, base=2)
-np.logspace(0, 2, base=2)
-np.logspace(0, 2, base=[1j, 2j], num=2)
-
-np.geomspace(1, 2)
diff --git a/numpy/typing/tests/data/pass/literal.py b/numpy/typing/tests/data/pass/literal.py
index 321ce3c2b..8eaeb6afb 100644
--- a/numpy/typing/tests/data/pass/literal.py
+++ b/numpy/typing/tests/data/pass/literal.py
@@ -31,6 +31,8 @@ order_list: List[Tuple[frozenset, Callable]] = [
     (KACF, partial(np.add, 1, 1)),  # i.e. np.ufunc.__call__
     (ACF, partial(np.reshape, AR, 1)),
     (KACF, partial(np.ravel, AR)),
+    (KACF, partial(np.asarray, 1)),
+    (KACF, partial(np.asanyarray, 1)),
 ]
 
 for order_set, func in order_list:
diff --git a/numpy/typing/tests/data/pass/simple.py b/numpy/typing/tests/data/pass/simple.py
index 527050557..4d397bd29 100644
--- a/numpy/typing/tests/data/pass/simple.py
+++ b/numpy/typing/tests/data/pass/simple.py
@@ -17,15 +17,6 @@ ndarray_func(np.array([1, 2]))
 array == 1
 array.dtype == float
 
-# Array creation routines checks
-np.array(1, dtype=float)
-np.array(1, copy=False)
-np.array(1, order='F')
-np.array(1, order=None)
-np.array(1, subok=True)
-np.array(1, ndmin=3)
-np.array(1, str, copy=True, order='C', subok=False, ndmin=2)
-
 ndarray_func(np.zeros([1, 2]))
 ndarray_func(np.ones([1, 2]))
 ndarray_func(np.empty([1, 2]))
diff --git a/numpy/typing/tests/data/reveal/array_constructors.py b/numpy/typing/tests/data/reveal/array_constructors.py
new file mode 100644
index 000000000..ba8a8eda1
--- /dev/null
+++ b/numpy/typing/tests/data/reveal/array_constructors.py
@@ -0,0 +1,42 @@
+from typing import List
+import numpy as np
+
+class SubClass(np.ndarray): ...
+
+A: np.ndarray
+B: SubClass
+C: List[int]
+
+reveal_type(np.asarray(A))  # E: ndarray
+reveal_type(np.asarray(B))  # E: ndarray
+reveal_type(np.asarray(C))  # E: ndarray
+
+reveal_type(np.asanyarray(A))  # E: ndarray
+reveal_type(np.asanyarray(B))  # E: SubClass
+reveal_type(np.asanyarray(B, dtype=int))  # E: ndarray
+reveal_type(np.asanyarray(C))  # E: ndarray
+
+reveal_type(np.ascontiguousarray(A))  # E: ndarray
+reveal_type(np.ascontiguousarray(B))  # E: ndarray
+reveal_type(np.ascontiguousarray(C))  # E: ndarray
+
+reveal_type(np.asfortranarray(A))  # E: ndarray
+reveal_type(np.asfortranarray(B))  # E: ndarray
+reveal_type(np.asfortranarray(C))  # E: ndarray
+
+reveal_type(np.require(A))  # E: ndarray
+reveal_type(np.require(B))  # E: SubClass
+reveal_type(np.require(B, requirements=None))  # E: SubClass
+reveal_type(np.require(B, dtype=int))  # E: ndarray
+reveal_type(np.require(B, requirements="E"))  # E: ndarray
+reveal_type(np.require(B, requirements=["ENSUREARRAY"]))  # E: ndarray
+reveal_type(np.require(B, requirements={"F", "E"}))  # E: ndarray
+reveal_type(np.require(B, requirements=["C", "OWNDATA"]))  # E: SubClass
+reveal_type(np.require(B, requirements="W"))  # E: SubClass
+reveal_type(np.require(B, requirements="A"))  # E: SubClass
+reveal_type(np.require(C))  # E: ndarray
+
+reveal_type(np.linspace(0, 10))  # E: numpy.ndarray
+reveal_type(np.linspace(0, 10, retstep=True))  # E: Tuple[numpy.ndarray, numpy.inexact]
+reveal_type(np.logspace(0, 10))  # E: numpy.ndarray
+reveal_type(np.geomspace(1, 10))  # E: numpy.ndarray
diff --git a/numpy/typing/tests/data/reveal/linspace.py b/numpy/typing/tests/data/reveal/linspace.py
deleted file mode 100644
index cfbbdf390..000000000
--- a/numpy/typing/tests/data/reveal/linspace.py
+++ /dev/null
@@ -1,6 +0,0 @@
-import numpy as np
-
-reveal_type(np.linspace(0, 10))  # E: numpy.ndarray
-reveal_type(np.linspace(0, 10, retstep=True))  # E: Tuple[numpy.ndarray, numpy.inexact]
-reveal_type(np.logspace(0, 10))  # E: numpy.ndarray
-reveal_type(np.geomspace(1, 10))  # E: numpy.ndarray