ENH: Add annotations for all remaining multiarray functions

2 files changed, 440 insertions, 45 deletions

diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
index 59e4e3694..6e24f6bff 100644
--- a/numpy/__init__.pyi
+++ b/numpy/__init__.pyi
@@ -326,6 +326,21 @@ from numpy.core.multiarray import (
     asanyarray as asanyarray,
     ascontiguousarray as ascontiguousarray,
     asfortranarray as asfortranarray,
+    arange as arange,
+    busday_count as busday_count,
+    busday_offset as busday_offset,
+    compare_chararrays as compare_chararrays,
+    datetime_as_string as datetime_as_string,
+    datetime_data as datetime_data,
+    frombuffer as frombuffer,
+    fromfile as fromfile,
+    fromiter as fromiter,
+    is_busday as is_busday,
+    promote_types as promote_types,
+    seterrobj as seterrobj,
+    geterrobj as geterrobj,
+    fromstring as fromstring,
+    frompyfunc as frompyfunc,
 )
 
 from numpy.core.numeric import (
@@ -611,10 +626,6 @@ class MachAr:
     ) -> None: ...
     def __getattr__(self, key: str) -> Any: ...
 
-class busdaycalendar:
-    def __new__(cls, weekmask: Any = ..., holidays: Any = ...) -> Any: ...
-    def __getattr__(self, key: str) -> Any: ...
-
 class chararray(ndarray[_ShapeType, _DType_co]):
     def __new__(
         subtype,
@@ -878,46 +889,8 @@ def round(a, decimals=..., out=...): ...
 def round_(a, decimals=..., out=...): ...
 def show_config(): ...
 
-# Placeholders for C-based functions
 # TODO: Sort out which parameters are positional-only
-@overload
-def arange(stop, dtype=..., *, like=...): ...
-@overload
-def arange(start, stop, step=..., dtype=..., *, like=...): ...
-def busday_count(
-    begindates,
-    enddates,
-    weekmask=...,
-    holidays=...,
-    busdaycal=...,
-    out=...,
-): ...
-def busday_offset(
-    dates,
-    offsets,
-    roll=...,
-    weekmask=...,
-    holidays=...,
-    busdaycal=...,
-    out=...,
-): ...
-def compare_chararrays(a, b, cmp_op, rstrip): ...
-def datetime_as_string(arr, unit=..., timezone=..., casting=...): ...
-def datetime_data(__dtype): ...
-def frombuffer(buffer, dtype=..., count=..., offset=..., *, like=...): ...
-def fromfile(
-    file, dtype=..., count=..., sep=..., offset=..., *, like=...
-): ...
-def fromiter(iter, dtype, count=..., *, like=...): ...
-def frompyfunc(func, nin, nout, * identity): ...
-def fromstring(string, dtype=..., count=..., sep=..., *, like=...): ...
-def geterrobj(): ...
-def is_busday(
-    dates, weekmask=..., holidays=..., busdaycal=..., out=...
-): ...
-def nested_iters(*args, **kwargs): ...  # TODO: Sort out parameters
-def promote_types(type1, type2): ...
-def seterrobj(errobj): ...
+def nested_iters(*args, **kwargs): ... # TODO: Sort out parameters
 
 _NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray)
 _DTypeScalar_co = TypeVar("_DTypeScalar_co", covariant=True, bound=generic)
@@ -3749,3 +3722,14 @@ class broadcast:
     def __next__(self) -> Tuple[Any, ...]: ...
     def __iter__(self: _T) -> _T: ...
     def reset(self) -> None: ...
+
+class busdaycalendar:
+    def __new__(
+        cls,
+        weekmask: ArrayLike = ...,
+        holidays: ArrayLike = ...,
+    ) -> busdaycalendar: ...
+    @property
+    def weekmask(self) -> NDArray[bool_]: ...
+    @property
+    def holidays(self) -> NDArray[datetime64]: ...
diff --git a/numpy/core/multiarray.pyi b/numpy/core/multiarray.pyi
index 30b3f416a..bc33165be 100644
--- a/numpy/core/multiarray.pyi
+++ b/numpy/core/multiarray.pyi
@@ -1,8 +1,13 @@
 # TODO: Sort out any and all missing functions in this namespace
 
+import os
 import sys
+import datetime as dt
 from typing import (
     Any,
+    Callable,
+    IO,
+    Iterable,
     Optional,
     overload,
     TypeVar,
@@ -14,15 +19,24 @@ from typing import (
 )
 
 from numpy import (
+    # Re-exports
     busdaycalendar as busdaycalendar,
-    ndarray,
-    dtype,
+    broadcast as broadcast,
+    dtype as dtype,
+    ndarray as ndarray,
+    nditer as nditer,
+
+    # The rest
+    nditer,
+    ufunc,
     str_,
     bool_,
     uint8,
     intp,
+    int_,
     float64,
     timedelta64,
+    datetime64,
     generic,
     unsignedinteger,
     signedinteger,
@@ -32,6 +46,7 @@ from numpy import (
     _OrderCF,
     _CastingKind,
     _ModeKind,
+    _SupportsBuffer,
 )
 
 from numpy.typing import (
@@ -53,8 +68,14 @@ from numpy.typing import (
     _ArrayLikeFloat_co,
     _ArrayLikeComplex_co,
     _ArrayLikeTD64_co,
+    _ArrayLikeDT64_co,
     _ArrayLikeObject_co,
+    _ArrayLikeStr_co,
+    _ArrayLikeBytes_co,
+    _ScalarLike_co,
     _IntLike_co,
+    _FloatLike_co,
+    _TD64Like_co,
 )
 
 if sys.version_info >= (3, 8):
@@ -65,6 +86,7 @@ else:
 _SCT = TypeVar("_SCT", bound=generic)
 _ArrayType = TypeVar("_ArrayType", bound=NDArray[Any])
 
+# Subscriptable subsets of `npt.DTypeLike` and `npt.ArrayLike`
 _DTypeLike = Union[
     dtype[_SCT],
     Type[_SCT],
@@ -72,6 +94,31 @@ _DTypeLike = Union[
 ]
 _ArrayLike = _NestedSequence[_SupportsArray[dtype[_SCT]]]
 
+# Valid time units
+_UnitKind = L[
+    "Y",
+    "M",
+    "D",
+    "h",
+    "m",
+    "s",
+    "ms",
+    "us", "μs",
+    "ns",
+    "ps",
+    "fs",
+    "as",
+]
+_RollKind = L[  # `raise` is deliberately excluded
+    "nat",
+    "forward",
+    "following",
+    "backward",
+    "preceding",
+    "modifiedfollowing",
+    "modifiedpreceding",
+]
+
 __all__: List[str]
 
 ALLOW_THREADS: Final[int]  # 0 or 1 (system-specific)
@@ -537,3 +584,367 @@ def asfortranarray(
     *,
     like: ArrayLike = ...,
 ) -> NDArray[Any]: ...
+
+# In practice `List[Any]` is list with an int, int and a valid
+# `np.seterrcall()` object
+def geterrobj() -> List[Any]: ...
+def seterrobj(__errobj: List[Any]) -> None: ...
+
+def promote_types(__type1: DTypeLike, __type2: DTypeLike) -> dtype[Any]: ...
+
+# `sep` is a de facto mandatory argument, as its default value is deprecated
+@overload
+def fromstring(
+    string: str | bytes,
+    dtype: None = ...,
+    count: SupportsIndex = ...,
+    *,
+    sep: str,
+    like: ArrayLike = ...,
+) -> NDArray[float64]: ...
+@overload
+def fromstring(
+    string: str | bytes,
+    dtype: _DTypeLike[_SCT],
+    count: SupportsIndex = ...,
+    *,
+    sep: str,
+    like: ArrayLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def fromstring(
+    string: str | bytes,
+    dtype: DTypeLike,
+    count: SupportsIndex = ...,
+    *,
+    sep: str,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
+
+def frompyfunc(
+    __func: Callable[..., Any],
+    nin: SupportsIndex,
+    nout: SupportsIndex,
+    *,
+    identity: Any = ...,
+) -> ufunc: ...
+
+@overload
+def fromfile(
+    file: str | bytes | os.PathLike[Any] | IO[Any],
+    dtype: None = ...,
+    count: SupportsIndex = ...,
+    sep: str = ...,
+    offset: SupportsIndex = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[float64]: ...
+@overload
+def fromfile(
+    file: str | bytes | os.PathLike[Any] | IO[Any],
+    dtype: _DTypeLike[_SCT],
+    count: SupportsIndex = ...,
+    sep: str = ...,
+    offset: SupportsIndex = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def fromfile(
+    file: str | bytes | os.PathLike[Any] | IO[Any],
+    dtype: DTypeLike,
+    count: SupportsIndex = ...,
+    sep: str = ...,
+    offset: SupportsIndex = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
+
+@overload
+def fromiter(
+    iter: Iterable[Any],
+    dtype: _DTypeLike[_SCT],
+    count: SupportsIndex = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def fromiter(
+    iter: Iterable[Any],
+    dtype: DTypeLike,
+    count: SupportsIndex = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
+
+@overload
+def frombuffer(
+    buffer: _SupportsBuffer,
+    dtype: None = ...,
+    count: SupportsIndex = ...,
+    offset: SupportsIndex = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[float64]: ...
+@overload
+def frombuffer(
+    buffer: _SupportsBuffer,
+    dtype: _DTypeLike[_SCT],
+    count: SupportsIndex = ...,
+    offset: SupportsIndex = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def frombuffer(
+    buffer: _SupportsBuffer,
+    dtype: DTypeLike,
+    count: SupportsIndex = ...,
+    offset: SupportsIndex = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
+
+@overload
+def arange(  # type: ignore[misc]
+    __stop: _IntLike_co,
+    *,
+    dtype: None = ...,
+    like: ArrayLike = ...,
+) -> NDArray[signedinteger[Any]]: ...
+@overload
+def arange(  # type: ignore[misc]
+    start: _IntLike_co,
+    stop: _IntLike_co,
+    step: _IntLike_co = ...,
+    dtype: None = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[signedinteger[Any]]: ...
+@overload
+def arange(  # type: ignore[misc]
+    __stop: _FloatLike_co,
+    *,
+    dtype: None = ...,
+    like: ArrayLike = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def arange(  # type: ignore[misc]
+    start: _FloatLike_co,
+    stop: _FloatLike_co,
+    step: _FloatLike_co = ...,
+    dtype: None = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[floating[Any]]: ...
+@overload
+def arange(
+    __stop: _TD64Like_co,
+    *,
+    dtype: None = ...,
+    like: ArrayLike = ...,
+) -> NDArray[timedelta64]: ...
+@overload
+def arange(
+    start: _TD64Like_co,
+    stop: _TD64Like_co,
+    step: _TD64Like_co = ...,
+    dtype: None = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[timedelta64]: ...
+@overload
+def arange(  # both start and stop must always be specified for datetime64
+    start: datetime64,
+    stop: datetime64,
+    step: datetime64 = ...,
+    dtype: None = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[datetime64]: ...
+@overload
+def arange(
+    __stop: Any,
+    *,
+    dtype: _DTypeLike[_SCT],
+    like: ArrayLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def arange(
+    start: Any,
+    stop: Any,
+    step: Any = ...,
+    dtype: _DTypeLike[_SCT] = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[_SCT]: ...
+@overload
+def arange(
+    __stop: Any,
+    *,
+    dtype: DTypeLike,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
+@overload
+def arange(
+    start: Any,
+    stop: Any,
+    step: Any = ...,
+    dtype: DTypeLike = ...,
+    *,
+    like: ArrayLike = ...,
+) -> NDArray[Any]: ...
+
+def datetime_data(
+    __dtype: str | _DTypeLike[datetime64] | _DTypeLike[timedelta64],
+) -> Tuple[str, int]: ...
+
+# The datetime functions perform unsafe casts to `datetime64[D]`,
+# so a lot of different argument types are allowed here
+
+@overload
+def busday_count(  # type: ignore[misc]
+    begindates: _ScalarLike_co,
+    enddates: _ScalarLike_co,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: None = ...,
+) -> int_: ...
+@overload
+def busday_count(  # type: ignore[misc]
+    begindates: ArrayLike,
+    enddates: ArrayLike,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: None = ...,
+) -> NDArray[int_]: ...
+@overload
+def busday_count(
+    begindates: ArrayLike,
+    enddates: ArrayLike,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: _ArrayType = ...,
+) -> _ArrayType: ...
+
+# `roll="raise"` is (more or less?) equivalent to `casting="safe"`
+@overload
+def busday_offset(  # type: ignore[misc]
+    dates: datetime64,
+    offsets: _TD64Like_co,
+    roll: L["raise"] = ...,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: None = ...,
+) -> datetime64: ...
+@overload
+def busday_offset(  # type: ignore[misc]
+    dates: _ArrayLike[datetime64],
+    offsets: _ArrayLikeTD64_co,
+    roll: L["raise"] = ...,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: None = ...,
+) -> NDArray[datetime64]: ...
+@overload
+def busday_offset(  # type: ignore[misc]
+    dates: _ArrayLike[datetime64],
+    offsets: _ArrayLike[timedelta64],
+    roll: L["raise"] = ...,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: _ArrayType = ...,
+) -> _ArrayType: ...
+@overload
+def busday_offset(  # type: ignore[misc]
+    dates: _ScalarLike_co,
+    offsets: _ScalarLike_co,
+    roll: _RollKind,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: None = ...,
+) -> datetime64: ...
+@overload
+def busday_offset(  # type: ignore[misc]
+    dates: ArrayLike,
+    offsets: ArrayLike,
+    roll: _RollKind,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: None = ...,
+) -> NDArray[datetime64]: ...
+@overload
+def busday_offset(
+    dates: ArrayLike,
+    offsets: ArrayLike,
+    roll: _RollKind,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: _ArrayType = ...,
+) -> _ArrayType: ...
+
+@overload
+def is_busday(  # type: ignore[misc]
+    dates: _ScalarLike_co,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: None = ...,
+) -> bool_: ...
+@overload
+def is_busday(  # type: ignore[misc]
+    dates: ArrayLike,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: None = ...,
+) -> NDArray[bool_]: ...
+@overload
+def is_busday(
+    dates: ArrayLike,
+    weekmask: ArrayLike = ...,
+    holidays: None | ArrayLike = ...,
+    busdaycal: None | busdaycalendar = ...,
+    out: _ArrayType = ...,
+) -> _ArrayType: ...
+
+@overload
+def datetime_as_string(  # type: ignore[misc]
+    arr: datetime64,
+    unit: None | L["auto"] | _UnitKind = ...,
+    timezone: L["naive", "UTC", "local"] | dt.tzinfo = ...,
+    casting: _CastingKind = ...,
+) -> str_: ...
+@overload
+def datetime_as_string(
+    arr: _ArrayLikeDT64_co,
+    unit: None | L["auto"] | _UnitKind = ...,
+    timezone: L["naive", "UTC", "local"] | dt.tzinfo = ...,
+    casting: _CastingKind = ...,
+) -> NDArray[str_]: ...
+
+@overload
+def compare_chararrays(
+    a1: _ArrayLikeStr_co,
+    a2: _ArrayLikeStr_co,
+    cmp: L["<", "<=", "==", ">=", ">", "!="],
+    rstrip: bool,
+) -> NDArray[bool_]: ...
+@overload
+def compare_chararrays(
+    a1: _ArrayLikeBytes_co,
+    a2: _ArrayLikeBytes_co,
+    cmp: L["<", "<=", "==", ">=", ">", "!="],
+    rstrip: bool,
+) -> NDArray[bool_]: ...
+
+def add_docstring(__obj: Callable[..., Any], __docstring: str) -> None: ...