Merge pull request #19237 from BvB93/multiarray2

ENH: Add annotations to np.core.multiarray part 3/4

numpy/__init__.pyi

@@ -1,6 +1,8 @@
 import builtins
 import os
 import sys
+import mmap
+import array as _array
 import datetime as dt
 from abc import abstractmethod
 from types import TracebackType
@@ -324,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 (
@@ -609,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,
@@ -876,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)
@@ -1626,7 +1601,18 @@ _DType_co = TypeVar("_DType_co", covariant=True, bound=dtype[Any])
 # have proper shape support
 _ShapeType = TypeVar("_ShapeType", bound=Any)
 _NumberType = TypeVar("_NumberType", bound=number[Any])
-_BufferType = Union[ndarray, bytes, bytearray, memoryview]
+
+# There is currently no exhaustive way to type the buffer protocol,
+# as it is implemented exclusivelly in the C API (python/typing#593)
+_SupportsBuffer = Union[
+    bytes,
+    bytearray,
+    memoryview,
+    _array.array[Any],
+    mmap.mmap,
+    NDArray[Any],
+    generic,
+]
 
 _T = TypeVar("_T")
 _T_co = TypeVar("_T_co", covariant=True)
@@ -1668,7 +1654,7 @@ class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]):
         cls: Type[_ArraySelf],
         shape: _ShapeLike,
         dtype: DTypeLike = ...,
-        buffer: _BufferType = ...,
+        buffer: _SupportsBuffer = ...,
         offset: int = ...,
         strides: _ShapeLike = ...,
         order: _OrderKACF = ...,
@@ -3736,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]: ...

numpy/core/_add_newdocs.py

@@ -1280,7 +1280,7 @@
 
 add_newdoc('numpy.core.multiarray', 'fromstring',
     """
-    fromstring(string, dtype=float, count=-1, sep='', *, like=None)
+    fromstring(string, dtype=float, count=-1, *, sep, like=None)
 
     A new 1-D array initialized from text data in a string.
 
@@ -1346,16 +1346,16 @@
 
 add_newdoc('numpy.core.multiarray', 'compare_chararrays',
     """
-    compare_chararrays(a, b, cmp_op, rstrip)
+    compare_chararrays(a1, a2, cmp, rstrip)
 
     Performs element-wise comparison of two string arrays using the
     comparison operator specified by `cmp_op`.
 
     Parameters
     ----------
-    a, b : array_like
+    a1, a2 : array_like
         Arrays to be compared.
-    cmp_op : {"<", "<=", "==", ">=", ">", "!="}
+    cmp : {"<", "<=", "==", ">=", ">", "!="}
         Type of comparison.
     rstrip : Boolean
         If True, the spaces at the end of Strings are removed before the comparison.
@@ -4475,7 +4475,7 @@
 
 add_newdoc('numpy.core.umath', 'frompyfunc',
     """
-    frompyfunc(func, nin, nout, *[, identity])
+    frompyfunc(func, /, nin, nout, *[, identity])
 
     Takes an arbitrary Python function and returns a NumPy ufunc.
 
@@ -4589,7 +4589,7 @@
 
 add_newdoc('numpy.core.umath', 'seterrobj',
     """
-    seterrobj(errobj)
+    seterrobj(errobj, /)
 
     Set the object that defines floating-point error handling.
 

numpy/core/multiarray.py

@@ -31,7 +31,7 @@
     'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast',
     'compare_chararrays', 'concatenate', 'copyto', 'correlate', 'correlate2',
     'count_nonzero', 'c_einsum', 'datetime_as_string', 'datetime_data',
-    'digitize', 'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype',
+    'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype',
     'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat',
     'frombuffer', 'fromfile', 'fromiter', 'fromstring', 'inner',
     'interp', 'interp_complex', 'is_busday', 'lexsort',