🐛✨ fix and improve _core._asarray

src/numpy-stubs/_core/_asarray.pyi

@@ -1,34 +1,118 @@
 from collections.abc import Iterable
-from typing import Any, Literal, TypeAlias, overload
+from typing import Any, Final, Literal as L, TypeAlias, overload
 from typing_extensions import TypeVar
 
-from numpy._typing import DTypeLike, NDArray, _SupportsArrayFunc
+import numpy as np
+from _numtype import ToBool_nd, ToBytes_nd, ToComplex128_nd, ToFloat64_nd, ToIntP_nd, ToObject_nd, ToStr_nd
+from numpy._typing import DTypeLike, NDArray, _ArrayLike, _DTypeLike, _SupportsArrayFunc as _Like
 
-_ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
+__all__ = ["require"]
 
-_Requirements: TypeAlias = Literal[
-    "C", "C_CONTIGUOUS", "CONTIGUOUS", "F", "F_CONTIGUOUS", "FORTRAN", "A", "ALIGNED", "W", "WRITEABLE", "O", "OWNDATA"
-]
-_E: TypeAlias = Literal["E", "ENSUREARRAY"]
-_RequirementsWithE: TypeAlias = _Requirements | _E
+###
+
+_ArrayT = TypeVar("_ArrayT", bound=NDArray[Any])
+_ScalarT = TypeVar("_ScalarT", bound=np.generic[Any])
+
+_Req: TypeAlias = L[
+    "C", "C_CONTIGUOUS", "CONTIGUOUS",
+    "F", "F_CONTIGUOUS", "FORTRAN",
+    "A", "ALIGNED",
+    "W", "WRITEABLE",
+    "O", "OWNDATA",
+]  # fmt: skip
+_ReqE: TypeAlias = L[_Req, "E", "ENSUREARRAY"]
+_ToReqs: TypeAlias = _Req | Iterable[_Req]
+_ToReqsE: TypeAlias = _ReqE | Iterable[_ReqE]
+
+###
+
+POSSIBLE_FLAGS: Final[dict[_ReqE, L["C", "F", "A", "W", "O", "E"]]]
 
 @overload
 def require(
-    a: _ArrayType,
-    dtype: None = ...,
-    requirements: _Requirements | Iterable[_Requirements] | None = ...,
+    a: _ArrayT,
+    dtype: None = None,
+    requirements: _ToReqs | None = None,
     *,
-    like: _SupportsArrayFunc = ...,
-) -> _ArrayType: ...
+    like: _Like | None = None,
+) -> _ArrayT: ...
 @overload
 def require(
-    a: object, dtype: DTypeLike = ..., requirements: _E | Iterable[_RequirementsWithE] = ..., *, like: _SupportsArrayFunc = ...
-) -> NDArray[Any]: ...
+    a: ToBool_nd,
+    dtype: None = None,
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[np.bool]: ...
+@overload
+def require(  # type: ignore[overload-overlap]
+    a: ToIntP_nd,
+    dtype: None = None,
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[np.intp]: ...
+@overload
+def require(  # type: ignore[overload-overlap]
+    a: ToFloat64_nd,
+    dtype: None = None,
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[np.float64]: ...
+@overload
+def require(  # type: ignore[overload-overlap]
+    a: ToComplex128_nd,
+    dtype: None = None,
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[np.complex128]: ...
+@overload
+def require(  # type: ignore[overload-overlap]
+    a: ToBytes_nd,
+    dtype: None = None,
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[np.bytes_]: ...
+@overload
+def require(  # type: ignore[overload-overlap]
+    a: ToStr_nd,
+    dtype: None = None,
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[np.str_]: ...
+@overload
+def require(
+    a: ToObject_nd,
+    dtype: None = None,
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[np.object_]: ...
+@overload
+def require(
+    a: _ArrayLike[_ScalarT],
+    dtype: None = None,
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[_ScalarT]: ...
+@overload
+def require(
+    a: object,
+    dtype: _DTypeLike[_ScalarT],
+    requirements: _ToReqsE | None = None,
+    *,
+    like: _Like | None = None,
+) -> NDArray[_ScalarT]: ...
 @overload
 def require(
     a: object,
-    dtype: DTypeLike = ...,
-    requirements: _Requirements | Iterable[_Requirements] | None = ...,
+    dtype: DTypeLike | None = None,
+    requirements: _ToReqsE | None = None,
     *,
-    like: _SupportsArrayFunc = ...,
+    like: _Like | None = None,
 ) -> NDArray[Any]: ...

test/static/accept/array_constructors.pyi

@@ -121,13 +121,13 @@ assert_type(np.require(A), npt.NDArray[np.float64])
 assert_type(np.require(B), MyArray[np.float64])
 assert_type(np.require(B, requirements=None), MyArray[np.float64])
 assert_type(np.require(B, dtype=int), npt.NDArray[Any])
-assert_type(np.require(B, requirements="E"), npt.NDArray[Any])
-assert_type(np.require(B, requirements=["ENSUREARRAY"]), npt.NDArray[Any])
-assert_type(np.require(B, requirements={"F", "E"}), npt.NDArray[Any])
+assert_type(np.require(B, requirements="E"), npt.NDArray[np.float64])
+assert_type(np.require(B, requirements=["ENSUREARRAY"]), npt.NDArray[np.float64])
+assert_type(np.require(B, requirements={"F", "E"}), npt.NDArray[np.float64])
 assert_type(np.require(B, requirements=["C", "OWNDATA"]), MyArray[np.float64])
 assert_type(np.require(B, requirements="W"), MyArray[np.float64])
 assert_type(np.require(B, requirements="A"), MyArray[np.float64])
-assert_type(np.require(C), npt.NDArray[Any])
+assert_type(np.require(C), npt.NDArray[np.intp])
 
 assert_type(np.linspace(0, 10), npt.NDArray[np.floating])
 assert_type(np.linspace(0, 10j), npt.NDArray[np.inexact])