From b4f22cf00138640baaa60fa127a2a50059c48152 Mon Sep 17 00:00:00 2001 From: jorenham Date: Sun, 28 Dec 2025 20:11:26 +0100 Subject: [PATCH] TYP: ``bincount`` shape-typing --- numpy/_core/multiarray.pyi | 15 ++++++++------- numpy/typing/tests/data/reveal/multiarray.pyi | 2 +- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/numpy/_core/multiarray.pyi b/numpy/_core/multiarray.pyi index e516c3cdab72..c6dfe3c2698e 100644 --- a/numpy/_core/multiarray.pyi +++ b/numpy/_core/multiarray.pyi @@ -213,6 +213,9 @@ type _ArangeScalar = np.integer | np.floating | np.datetime64 | np.timedelta64 type _ToDates = dt.date | _NestedSequence[dt.date] type _ToDeltas = dt.timedelta | _NestedSequence[dt.timedelta] +type _BitOrder = L["big", "little"] +type _MaxWork = L[-1, 0] + @type_check_only class _SupportsArray[ArrayT_co: np.ndarray](Protocol): def __array__(self, /) -> ArrayT_co: ... @@ -619,15 +622,15 @@ def vdot(a: _ArrayLikeObject_co, b: object, /) -> Any: ... @overload def vdot(a: object, b: _ArrayLikeObject_co, /) -> Any: ... -def bincount(x: ArrayLike, /, weights: ArrayLike | None = None, minlength: SupportsIndex = 0) -> NDArray[intp]: ... +# +def bincount(x: ArrayLike, /, weights: ArrayLike | None = None, minlength: SupportsIndex = 0) -> _Array1D[intp]: ... +# def copyto(dst: ndarray, src: ArrayLike, casting: _CastingKind = "same_kind", where: object = True) -> None: ... def putmask(a: ndarray, /, mask: _ArrayLikeBool_co, values: ArrayLike) -> None: ... -type _BitOrder = L["big", "little"] - @overload -def packbits(a: _ArrayLikeInt_co, /, axis: None = None, bitorder: _BitOrder = "big") -> ndarray[tuple[int], dtype[uint8]]: ... +def packbits(a: _ArrayLikeInt_co, /, axis: None = None, bitorder: _BitOrder = "big") -> _Array1D[uint8]: ... @overload def packbits(a: _ArrayLikeInt_co, /, axis: SupportsIndex, bitorder: _BitOrder = "big") -> NDArray[uint8]: ... @@ -638,7 +641,7 @@ def unpackbits( axis: None = None, count: SupportsIndex | None = None, bitorder: _BitOrder = "big", -) -> ndarray[tuple[int], dtype[uint8]]: ... +) -> _Array1D[uint8]: ... @overload def unpackbits( a: _ArrayLike[uint8], @@ -648,8 +651,6 @@ def unpackbits( bitorder: _BitOrder = "big", ) -> NDArray[uint8]: ... -type _MaxWork = L[-1, 0] - # any two python objects will be accepted, not just `ndarray`s def shares_memory(a: object, b: object, /, max_work: _MaxWork = -1) -> bool: ... def may_share_memory(a: object, b: object, /, max_work: _MaxWork = 0) -> bool: ... diff --git a/numpy/typing/tests/data/reveal/multiarray.pyi b/numpy/typing/tests/data/reveal/multiarray.pyi index ada8f7777696..fefb6e2fbb5d 100644 --- a/numpy/typing/tests/data/reveal/multiarray.pyi +++ b/numpy/typing/tests/data/reveal/multiarray.pyi @@ -86,7 +86,7 @@ assert_type(np.vdot(AR_LIKE_f, AR_i8), np.floating) assert_type(np.vdot(AR_u1, 1), np.signedinteger) assert_type(np.vdot(1.5j, 1), np.complexfloating) -assert_type(np.bincount(AR_i8), npt.NDArray[np.intp]) +assert_type(np.bincount(AR_i8), np.ndarray[tuple[int], np.dtype[np.intp]]) assert_type(np.copyto(AR_f8, [1., 1.5, 1.6]), None)