ENH: Add annotations for the `np.rec` functions

1 files changed, 169 insertions, 51 deletions

diff --git a/numpy/core/records.pyi b/numpy/core/records.pyi
index 198636058..fda118276 100644
--- a/numpy/core/records.pyi
+++ b/numpy/core/records.pyi
@@ -1,65 +1,183 @@
-from typing import List
+import os
+from typing import (
+    List,
+    Sequence,
+    Any,
+    TypeVar,
+    Iterable,
+    overload,
+    Tuple,
+    Protocol,
+)
 
 from numpy import (
     format_parser as format_parser,
     record as record,
     recarray as recarray,
+    dtype,
+    generic,
+    void,
+    _ByteOrder,
+    _SupportsBuffer,
+)
+
+from numpy.typing import (
+    ArrayLike,
+    DTypeLike,
+    NDArray,
+    _ShapeLike,
+    _ArrayLikeVoid_co,
+    _NestedSequence,
 )
 
+_SCT = TypeVar("_SCT", bound=generic)
+
+_RecArray = recarray[Any, dtype[_SCT]]
+
+class _SupportsReadInto(Protocol):
+    def seek(self, offset: int, whence: int, /) -> object: ...
+    def tell(self, /) -> int: ...
+    def readinto(self, buffer: memoryview, /) -> int: ...
+
 __all__: List[str]
 
+@overload
 def fromarrays(
-    arrayList,
-    dtype=...,
-    shape=...,
-    formats=...,
-    names=...,
-    titles=...,
-    aligned=...,
-    byteorder=...,
-): ...
+    arrayList: Iterable[ArrayLike],
+    dtype: DTypeLike = ...,
+    shape: None | _ShapeLike = ...,
+    formats: None = ...,
+    names: None = ...,
+    titles: None = ...,
+    aligned: bool = ...,
+    byteorder: None = ...,
+) -> _RecArray[Any]: ...
+@overload
+def fromarrays(
+    arrayList: Iterable[ArrayLike],
+    dtype: None = ...,
+    shape: None | _ShapeLike = ...,
+    *,
+    formats: DTypeLike,
+    names: None | str | Sequence[str] = ...,
+    titles: None | str | Sequence[str] = ...,
+    aligned: bool = ...,
+    byteorder: None | _ByteOrder = ...,
+) -> _RecArray[record]: ...
+
+@overload
 def fromrecords(
-    recList,
-    dtype=...,
-    shape=...,
-    formats=...,
-    names=...,
-    titles=...,
-    aligned=...,
-    byteorder=...,
-): ...
+    recList: _ArrayLikeVoid_co | Tuple[Any, ...] | _NestedSequence[Tuple[Any, ...]],
+    dtype: DTypeLike = ...,
+    shape: None | _ShapeLike = ...,
+    formats: None = ...,
+    names: None = ...,
+    titles: None = ...,
+    aligned: bool = ...,
+    byteorder: None = ...,
+) -> _RecArray[record]: ...
+@overload
+def fromrecords(
+    recList: _ArrayLikeVoid_co | Tuple[Any, ...] | _NestedSequence[Tuple[Any, ...]],
+    dtype: None = ...,
+    shape: None | _ShapeLike = ...,
+    *,
+    formats: DTypeLike,
+    names: None | str | Sequence[str] = ...,
+    titles: None | str | Sequence[str] = ...,
+    aligned: bool = ...,
+    byteorder: None | _ByteOrder = ...,
+) -> _RecArray[record]: ...
+
+@overload
 def fromstring(
-    datastring,
-    dtype=...,
-    shape=...,
-    offset=...,
-    formats=...,
-    names=...,
-    titles=...,
-    aligned=...,
-    byteorder=...,
-): ...
+    datastring: _SupportsBuffer,
+    dtype: DTypeLike,
+    shape: None | _ShapeLike = ...,
+    offset: int = ...,
+    formats: None = ...,
+    names: None = ...,
+    titles: None = ...,
+    aligned: bool = ...,
+    byteorder: None = ...,
+) -> _RecArray[record]: ...
+@overload
+def fromstring(
+    datastring: _SupportsBuffer,
+    dtype: None = ...,
+    shape: None | _ShapeLike = ...,
+    offset: int = ...,
+    *,
+    formats: DTypeLike,
+    names: None | str | Sequence[str] = ...,
+    titles: None | str | Sequence[str] = ...,
+    aligned: bool = ...,
+    byteorder: None | _ByteOrder = ...,
+) -> _RecArray[record]: ...
+
+@overload
 def fromfile(
-    fd,
-    dtype=...,
-    shape=...,
-    offset=...,
-    formats=...,
-    names=...,
-    titles=...,
-    aligned=...,
-    byteorder=...,
-): ...
+    fd: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _SupportsReadInto,
+    dtype: DTypeLike,
+    shape: None | _ShapeLike = ...,
+    offset: int = ...,
+    formats: None = ...,
+    names: None = ...,
+    titles: None = ...,
+    aligned: bool = ...,
+    byteorder: None = ...,
+) -> _RecArray[Any]: ...
+@overload
+def fromfile(
+    fd: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _SupportsReadInto,
+    dtype: None = ...,
+    shape: None | _ShapeLike = ...,
+    offset: int = ...,
+    *,
+    formats: DTypeLike,
+    names: None | str | Sequence[str] = ...,
+    titles: None | str | Sequence[str] = ...,
+    aligned: bool = ...,
+    byteorder: None | _ByteOrder = ...,
+) -> _RecArray[record]: ...
+
+@overload
+def array(
+    obj: _SCT | NDArray[_SCT],
+    dtype: None = ...,
+    shape: None | _ShapeLike = ...,
+    offset: int = ...,
+    formats: None = ...,
+    names: None = ...,
+    titles: None = ...,
+    aligned: bool = ...,
+    byteorder: None = ...,
+    copy: bool = ...,
+) -> _RecArray[_SCT]: ...
+@overload
+def array(
+    obj: ArrayLike,
+    dtype: DTypeLike,
+    shape: None | _ShapeLike = ...,
+    offset: int = ...,
+    formats: None = ...,
+    names: None = ...,
+    titles: None = ...,
+    aligned: bool = ...,
+    byteorder: None = ...,
+    copy: bool = ...,
+) -> _RecArray[Any]: ...
+@overload
 def array(
-    obj,
-    dtype=...,
-    shape=...,
-    offset=...,
-    strides=...,
-    formats=...,
-    names=...,
-    titles=...,
-    aligned=...,
-    byteorder=...,
-    copy=...,
-): ...
+    obj: ArrayLike,
+    dtype: None = ...,
+    shape: None | _ShapeLike = ...,
+    offset: int = ...,
+    *,
+    formats: DTypeLike,
+    names: None | str | Sequence[str] = ...,
+    titles: None | str | Sequence[str] = ...,
+    aligned: bool = ...,
+    byteorder: None | _ByteOrder = ...,
+    copy: bool = ...,
+) -> _RecArray[record]: ...