diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a4da3323..ae47e943 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -57,6 +57,7 @@ jobs: src/numpy-stubs/lib/__init__.pyi src/numpy-stubs/lib/_array_utils_impl.pyi src/numpy-stubs/lib/_arraypad_impl.pyi + src/numpy-stubs/lib/_arraysetops_impl.pyi src/numpy-stubs/lib/_datasource.pyi src/numpy-stubs/lib/_histograms_impl.pyi src/numpy-stubs/lib/_iotools.pyi @@ -110,6 +111,7 @@ jobs: src/numpy-stubs/lib/__init__.pyi src/numpy-stubs/lib/_array_utils_impl.pyi src/numpy-stubs/lib/_arraypad_impl.pyi + src/numpy-stubs/lib/_arraysetops_impl.pyi src/numpy-stubs/lib/_datasource.pyi src/numpy-stubs/lib/_histograms_impl.pyi src/numpy-stubs/lib/_iotools.pyi diff --git a/src/numpy-stubs/lib/_arraysetops_impl.pyi b/src/numpy-stubs/lib/_arraysetops_impl.pyi index 236a44f9..9bd51b0e 100644 --- a/src/numpy-stubs/lib/_arraysetops_impl.pyi +++ b/src/numpy-stubs/lib/_arraysetops_impl.pyi @@ -1,17 +1,26 @@ -from typing import Any, Generic, Literal as L, NamedTuple, SupportsIndex, overload +from typing import Any, Generic, Literal as L, NamedTuple, SupportsIndex as CanIndex, TypeAlias, overload from typing_extensions import TypeVar, deprecated import numpy as np -from numpy._typing import ( - ArrayLike, - NDArray, - _ArrayLike, - _ArrayLikeBool_co, - _ArrayLikeDT64_co, - _ArrayLikeNumber_co, - _ArrayLikeObject_co, - _ArrayLikeTD64_co, +from _numtype import ( + Array, + Array_1d, + CoBool_nd, + CoComplex128_nd, + CoComplex_nd, + CoDateTime_nd, + CoFloat64_nd, + CoFloating_nd, + CoIntP_nd, + CoInteger_nd, + CoTimeDelta_nd, + ToComplex128_nd, + ToFloat64_nd, + ToIntP_nd, + ToObject_nd, + _ToArray1_nd, ) +from numpy._typing import ArrayLike, _ArrayLike __all__ = [ "ediff1d", @@ -28,15 +37,9 @@ __all__ = [ "unique_values", ] -_SCT = TypeVar("_SCT", bound=np.generic) -_NumberT = TypeVar("_NumberT", bound=np.number) +_ScalarT = TypeVar("_ScalarT", bound=np.generic, default=Any) +_CoNumberT = TypeVar("_CoNumberT", bound=np.number | np.timedelta64 | np.object_) -# Explicitly set all allowed values to prevent accidental castings to -# abstract dtypes (their common super-type). -# -# Only relevant if two or more arguments are parametrized, (e.g. `setdiff1d`) -# which could result in, for example, `int64` and `float64`producing a -# `number[_64Bit]` array _AnyScalarT = TypeVar( "_AnyScalarT", np.bool, @@ -63,288 +66,594 @@ _AnyScalarT = TypeVar( np.void, ) -class UniqueAllResult(NamedTuple, Generic[_SCT]): - values: NDArray[_SCT] - indices: NDArray[np.int_] - inverse_indices: NDArray[np.int_] - counts: NDArray[np.int_] +_IntersectResult: TypeAlias = tuple[Array_1d[_ScalarT], Array_1d[np.intp], Array_1d[np.intp]] + +### + +class UniqueAllResult(NamedTuple, Generic[_ScalarT]): + values: Array_1d[_ScalarT] + indices: Array_1d[np.intp] + inverse_indices: Array[np.intp] + counts: Array_1d[np.intp] -class UniqueCountsResult(NamedTuple, Generic[_SCT]): - values: NDArray[_SCT] - counts: NDArray[np.int_] +class UniqueCountsResult(NamedTuple, Generic[_ScalarT]): + values: Array_1d[_ScalarT] + counts: Array_1d[np.intp] -class UniqueInverseResult(NamedTuple, Generic[_SCT]): - values: NDArray[_SCT] - inverse_indices: NDArray[np.int_] +class UniqueInverseResult(NamedTuple, Generic[_ScalarT]): + values: Array_1d[_ScalarT] + inverse_indices: Array[np.intp] +# @overload -def ediff1d(ary: _ArrayLikeBool_co, to_end: ArrayLike | None = ..., to_begin: ArrayLike | None = ...) -> NDArray[np.int8]: ... +def ediff1d( + ary: CoBool_nd, + to_end: ArrayLike | None = None, + to_begin: ArrayLike | None = None, +) -> Array_1d[np.int8]: ... +@overload +def ediff1d( # type: ignore[overload-overlap] + ary: ToIntP_nd, + to_end: CoInteger_nd | None = None, + to_begin: CoInteger_nd | None = None, +) -> Array_1d[np.intp]: ... +@overload +def ediff1d( # type: ignore[overload-overlap] + ary: ToFloat64_nd, + to_end: CoFloating_nd | None = None, + to_begin: CoFloating_nd | None = None, +) -> Array_1d[np.float64]: ... +@overload +def ediff1d( # type: ignore[overload-overlap] + ary: ToComplex128_nd, + to_end: CoComplex_nd | None = None, + to_begin: CoComplex_nd | None = None, +) -> Array_1d[np.complex128]: ... @overload -def ediff1d(ary: _ArrayLike[_NumberT], to_end: ArrayLike | None = ..., to_begin: ArrayLike | None = ...) -> NDArray[_NumberT]: ... +def ediff1d( + ary: CoDateTime_nd, + to_end: CoTimeDelta_nd | None = None, + to_begin: CoTimeDelta_nd | None = None, +) -> Array_1d[np.timedelta64]: ... @overload def ediff1d( - ary: _ArrayLikeDT64_co | _ArrayLikeTD64_co, - to_end: ArrayLike | None = ..., - to_begin: ArrayLike | None = ..., -) -> NDArray[np.timedelta64]: ... + ary: ToObject_nd, + to_end: ArrayLike | None = None, + to_begin: ArrayLike | None = None, +) -> Array_1d[np.object_]: ... @overload def ediff1d( - ary: _ArrayLikeObject_co, - to_end: ArrayLike | None = ..., - to_begin: ArrayLike | None = ..., -) -> NDArray[np.object_]: ... + ary: _ToArray1_nd[_CoNumberT], + to_end: ArrayLike | None = None, + to_begin: ArrayLike | None = None, +) -> Array_1d[_CoNumberT]: ... @overload -def ediff1d(ary: _ArrayLikeNumber_co, to_end: ArrayLike | None = ..., to_begin: ArrayLike | None = ...) -> NDArray[Any]: ... +def ediff1d( + ary: CoComplex_nd | CoTimeDelta_nd, + to_end: CoComplex_nd | CoTimeDelta_nd | None = None, + to_begin: CoComplex_nd | CoTimeDelta_nd | None = None, +) -> Array_1d[Any]: ... # @overload def unique( - ar: _ArrayLike[_SCT], - return_index: L[False] = ..., - return_inverse: L[False] = ..., - return_counts: L[False] = ..., - axis: SupportsIndex | None = ..., + ar: _ArrayLike[_ScalarT], + return_index: L[False] = False, + return_inverse: L[False] = False, + return_counts: L[False] = False, + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> NDArray[_SCT]: ... + equal_nan: bool = True, +) -> Array[_ScalarT]: ... @overload def unique( ar: ArrayLike, - return_index: L[False] = ..., - return_inverse: L[False] = ..., - return_counts: L[False] = ..., - axis: SupportsIndex | None = ..., + return_index: L[False] = False, + return_inverse: L[False] = False, + return_counts: L[False] = False, + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> NDArray[Any]: ... + equal_nan: bool = True, +) -> Array[Any]: ... @overload def unique( - ar: _ArrayLike[_SCT], - return_index: L[True] = ..., - return_inverse: L[False] = ..., - return_counts: L[False] = ..., - axis: SupportsIndex | None = ..., + ar: _ArrayLike[_ScalarT], + return_index: L[True], + return_inverse: L[False] = False, + return_counts: L[False] = False, + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[_SCT], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp]]: ... @overload def unique( ar: ArrayLike, - return_index: L[True] = ..., - return_inverse: L[False] = ..., - return_counts: L[False] = ..., - axis: SupportsIndex | None = ..., + return_index: L[True], + return_inverse: L[False] = False, + return_counts: L[False] = False, + axis: CanIndex | None = None, + *, + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp]]: ... +@overload +def unique( + ar: _ArrayLike[_ScalarT], + return_index: L[False], + return_inverse: L[True], + return_counts: L[False] = False, + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[Any], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp]]: ... @overload def unique( - ar: _ArrayLike[_SCT], - return_index: L[False] = ..., - return_inverse: L[True] = ..., - return_counts: L[False] = ..., - axis: SupportsIndex | None = ..., + ar: _ArrayLike[_ScalarT], + return_index: L[False] = False, + *, + return_inverse: L[True], + return_counts: L[False] = False, + axis: CanIndex | None = None, + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[False], + return_inverse: L[True], + return_counts: L[False] = False, + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[_SCT], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp]]: ... @overload def unique( ar: ArrayLike, - return_index: L[False] = ..., - return_inverse: L[True] = ..., - return_counts: L[False] = ..., - axis: SupportsIndex | None = ..., + return_index: L[False] = False, *, - equal_nan: bool = ..., -) -> tuple[NDArray[Any], NDArray[np.int_]]: ... + return_inverse: L[True], + return_counts: L[False] = False, + axis: CanIndex | None = None, + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp]]: ... @overload def unique( - ar: _ArrayLike[_SCT], - return_index: L[False] = ..., - return_inverse: L[False] = ..., - return_counts: L[True] = ..., - axis: SupportsIndex | None = ..., + ar: _ArrayLike[_ScalarT], + return_index: L[False], + return_inverse: L[False], + return_counts: L[True], + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[_SCT], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp]]: ... +@overload +def unique( + ar: _ArrayLike[_ScalarT], + return_index: L[False] = False, + return_inverse: L[False] = False, + *, + return_counts: L[True], + axis: CanIndex | None = None, + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp]]: ... @overload def unique( ar: ArrayLike, - return_index: L[False] = ..., - return_inverse: L[False] = ..., - return_counts: L[True] = ..., - axis: SupportsIndex | None = ..., + return_index: L[False], + return_inverse: L[False], + return_counts: L[True], + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[Any], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp]]: ... @overload def unique( - ar: _ArrayLike[_SCT], - return_index: L[True] = ..., - return_inverse: L[True] = ..., - return_counts: L[False] = ..., - axis: SupportsIndex | None = ..., + ar: ArrayLike, + return_index: L[False] = False, + return_inverse: L[False] = False, *, - equal_nan: bool = ..., -) -> tuple[NDArray[_SCT], NDArray[np.int_], NDArray[np.int_]]: ... + return_counts: L[True], + axis: CanIndex | None = None, + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp]]: ... +@overload +def unique( + ar: _ArrayLike[_ScalarT], + return_index: L[True], + return_inverse: L[True], + return_counts: L[False] = False, + axis: CanIndex | None = None, + *, + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp], Array[np.intp]]: ... @overload def unique( ar: ArrayLike, - return_index: L[True] = ..., - return_inverse: L[True] = ..., - return_counts: L[False] = ..., - axis: SupportsIndex | None = ..., + return_index: L[True], + return_inverse: L[True], + return_counts: L[False] = False, + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[Any], NDArray[np.int_], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp], Array[np.intp]]: ... @overload def unique( - ar: _ArrayLike[_SCT], - return_index: L[True] = ..., - return_inverse: L[False] = ..., - return_counts: L[True] = ..., - axis: SupportsIndex | None = ..., + ar: _ArrayLike[_ScalarT], + return_index: L[True], + return_inverse: L[False], + return_counts: L[True], + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[_SCT], NDArray[np.int_], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp], Array[np.intp]]: ... +@overload +def unique( + ar: _ArrayLike[_ScalarT], + return_index: L[True], + *, + return_inverse: L[False] = False, + return_counts: L[True], + axis: CanIndex | None = None, + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp], Array[np.intp]]: ... @overload def unique( ar: ArrayLike, - return_index: L[True] = ..., - return_inverse: L[False] = ..., - return_counts: L[True] = ..., - axis: SupportsIndex | None = ..., + return_index: L[True], + return_inverse: L[False], + return_counts: L[True], + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[Any], NDArray[np.int_], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp], Array[np.intp]]: ... @overload def unique( - ar: _ArrayLike[_SCT], - return_index: L[False] = ..., - return_inverse: L[True] = ..., - return_counts: L[True] = ..., - axis: SupportsIndex | None = ..., + ar: ArrayLike, + return_index: L[True], + return_inverse: L[False] = False, + *, + return_counts: L[True], + axis: CanIndex | None = None, + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp], Array[np.intp]]: ... +@overload +def unique( + ar: _ArrayLike[_ScalarT], + return_index: L[False], + return_inverse: L[True], + return_counts: L[True], + axis: CanIndex | None = None, + *, + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp], Array[np.intp]]: ... +@overload +def unique( + ar: _ArrayLike[_ScalarT], + return_index: L[False] = False, + *, + return_inverse: L[True], + return_counts: L[True], + axis: CanIndex | None = None, + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp], Array[np.intp]]: ... +@overload +def unique( + ar: ArrayLike, + return_index: L[False], + return_inverse: L[True], + return_counts: L[True], + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[_SCT], NDArray[np.int_], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp], Array[np.intp]]: ... @overload def unique( ar: ArrayLike, - return_index: L[False] = ..., - return_inverse: L[True] = ..., - return_counts: L[True] = ..., - axis: SupportsIndex | None = ..., + return_index: L[False] = False, *, - equal_nan: bool = ..., -) -> tuple[NDArray[Any], NDArray[np.int_], NDArray[np.int_]]: ... + return_inverse: L[True], + return_counts: L[True], + axis: CanIndex | None = None, + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp], Array[np.intp]]: ... @overload def unique( - ar: _ArrayLike[_SCT], - return_index: L[True] = ..., - return_inverse: L[True] = ..., - return_counts: L[True] = ..., - axis: SupportsIndex | None = ..., + ar: _ArrayLike[_ScalarT], + return_index: L[True], + return_inverse: L[True], + return_counts: L[True], + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[_SCT], NDArray[np.int_], NDArray[np.int_], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[_ScalarT], Array[np.intp], Array[np.intp], Array[np.intp]]: ... @overload def unique( ar: ArrayLike, - return_index: L[True] = ..., - return_inverse: L[True] = ..., - return_counts: L[True] = ..., - axis: SupportsIndex | None = ..., + return_index: L[True], + return_inverse: L[True], + return_counts: L[True], + axis: CanIndex | None = None, *, - equal_nan: bool = ..., -) -> tuple[NDArray[Any], NDArray[np.int_], NDArray[np.int_], NDArray[np.int_]]: ... + equal_nan: bool = True, +) -> tuple[Array[Any], Array[np.intp], Array[np.intp], Array[np.intp]]: ... # @overload -def unique_all(x: _ArrayLike[_SCT], /) -> UniqueAllResult[_SCT]: ... +def unique_all(x: CoBool_nd, /) -> UniqueAllResult[np.bool]: ... @overload -def unique_all(x: ArrayLike, /) -> UniqueAllResult[Any]: ... +def unique_all(x: ToIntP_nd, /) -> UniqueAllResult[np.intp]: ... +@overload +def unique_all(x: ToFloat64_nd, /) -> UniqueAllResult[np.float64]: ... +@overload +def unique_all(x: ToComplex128_nd, /) -> UniqueAllResult[np.complex128]: ... +@overload +def unique_all(x: _ArrayLike[_ScalarT], /) -> UniqueAllResult[_ScalarT]: ... +@overload +def unique_all(x: ArrayLike, /) -> UniqueAllResult: ... # @overload -def unique_counts(x: _ArrayLike[_SCT], /) -> UniqueCountsResult[_SCT]: ... +def unique_counts(x: CoBool_nd, /) -> UniqueCountsResult[np.bool]: ... +@overload +def unique_counts(x: ToIntP_nd, /) -> UniqueCountsResult[np.intp]: ... @overload -def unique_counts(x: ArrayLike, /) -> UniqueCountsResult[Any]: ... +def unique_counts(x: ToFloat64_nd, /) -> UniqueCountsResult[np.float64]: ... +@overload +def unique_counts(x: ToComplex128_nd, /) -> UniqueCountsResult[np.complex128]: ... +@overload +def unique_counts(x: _ArrayLike[_ScalarT], /) -> UniqueCountsResult[_ScalarT]: ... +@overload +def unique_counts(x: ArrayLike, /) -> UniqueCountsResult: ... # @overload -def unique_inverse(x: _ArrayLike[_SCT], /) -> UniqueInverseResult[_SCT]: ... +def unique_inverse(x: CoBool_nd, /) -> UniqueInverseResult[np.bool]: ... +@overload +def unique_inverse(x: ToIntP_nd, /) -> UniqueInverseResult[np.intp]: ... @overload -def unique_inverse(x: ArrayLike, /) -> UniqueInverseResult[Any]: ... +def unique_inverse(x: ToFloat64_nd, /) -> UniqueInverseResult[np.float64]: ... +@overload +def unique_inverse(x: ToComplex128_nd, /) -> UniqueInverseResult[np.complex128]: ... +@overload +def unique_inverse(x: _ArrayLike[_ScalarT], /) -> UniqueInverseResult[_ScalarT]: ... +@overload +def unique_inverse(x: ArrayLike, /) -> UniqueInverseResult: ... # @overload -def unique_values(x: _ArrayLike[_SCT], /) -> NDArray[_SCT]: ... +def unique_values(x: CoBool_nd, /) -> Array_1d[np.bool]: ... +@overload +def unique_values(x: ToIntP_nd, /) -> Array_1d[np.intp]: ... @overload -def unique_values(x: ArrayLike, /) -> NDArray[Any]: ... +def unique_values(x: ToFloat64_nd, /) -> Array_1d[np.float64]: ... +@overload +def unique_values(x: ToComplex128_nd, /) -> Array_1d[np.complex128]: ... +@overload +def unique_values(x: _ArrayLike[_ScalarT], /) -> Array_1d[_ScalarT]: ... +@overload +def unique_values(x: ArrayLike, /) -> Array_1d[Any]: ... # @overload +def intersect1d( + ar1: ToFloat64_nd, + ar2: CoFloat64_nd, + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[np.float64]: ... +@overload +def intersect1d( + ar1: ToFloat64_nd, + ar2: CoFloat64_nd, + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult[np.float64]: ... +@overload +def intersect1d( + ar1: CoFloat64_nd, + ar2: ToFloat64_nd, + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[np.float64]: ... +@overload +def intersect1d( + ar1: CoFloat64_nd, + ar2: ToFloat64_nd, + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult[np.float64]: ... +@overload +def intersect1d( + ar1: CoBool_nd, + ar2: CoBool_nd, + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[np.bool]: ... +@overload +def intersect1d( + ar1: CoBool_nd, + ar2: CoBool_nd, + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult[np.bool]: ... +@overload +def intersect1d( + ar1: ToIntP_nd, + ar2: CoIntP_nd, + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[np.intp]: ... +@overload +def intersect1d( + ar1: ToIntP_nd, + ar2: CoIntP_nd, + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult[np.intp]: ... +@overload +def intersect1d( + ar1: CoIntP_nd, + ar2: ToIntP_nd, + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[np.intp]: ... +@overload +def intersect1d( + ar1: CoIntP_nd, + ar2: ToIntP_nd, + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult[np.intp]: ... +@overload +def intersect1d( + ar1: ToComplex128_nd, + ar2: CoComplex128_nd, + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[np.complex128]: ... +@overload +def intersect1d( + ar1: ToComplex128_nd, + ar2: CoComplex128_nd, + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult[np.complex128]: ... +@overload +def intersect1d( + ar1: CoComplex128_nd, + ar2: ToComplex128_nd, + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[np.complex128]: ... +@overload +def intersect1d( + ar1: CoComplex128_nd, + ar2: ToComplex128_nd, + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult[np.complex128]: ... +@overload def intersect1d( ar1: _ArrayLike[_AnyScalarT], ar2: _ArrayLike[_AnyScalarT], - assume_unique: bool = ..., - return_indices: L[False] = ..., -) -> NDArray[_AnyScalarT]: ... + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[_AnyScalarT]: ... +@overload +def intersect1d( + ar1: _ArrayLike[_AnyScalarT], + ar2: _ArrayLike[_AnyScalarT], + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult[_AnyScalarT]: ... @overload def intersect1d( ar1: ArrayLike, ar2: ArrayLike, - assume_unique: bool = ..., - return_indices: L[False] = ..., -) -> NDArray[Any]: ... + assume_unique: bool = False, + return_indices: L[False] = False, +) -> Array_1d[Any]: ... @overload def intersect1d( + ar1: ArrayLike, + ar2: ArrayLike, + assume_unique: bool = False, + *, + return_indices: L[True], +) -> _IntersectResult: ... + +# +@overload +def union1d(ar1: ToFloat64_nd, ar2: CoFloat64_nd, assume_unique: bool = False) -> Array_1d[np.float64]: ... +@overload +def union1d(ar1: CoFloat64_nd, ar2: ToFloat64_nd, assume_unique: bool = False) -> Array_1d[np.float64]: ... +@overload +def union1d(ar1: CoBool_nd, ar2: CoBool_nd, assume_unique: bool = False) -> Array_1d[np.bool]: ... +@overload +def union1d(ar1: ToIntP_nd, ar2: CoIntP_nd, assume_unique: bool = False) -> Array_1d[np.intp]: ... +@overload +def union1d(ar1: CoIntP_nd, ar2: ToIntP_nd, assume_unique: bool = False) -> Array_1d[np.intp]: ... +@overload +def union1d(ar1: ToComplex128_nd, ar2: CoComplex128_nd, assume_unique: bool = False) -> Array_1d[np.complex128]: ... +@overload +def union1d(ar1: CoComplex128_nd, ar2: ToComplex128_nd, assume_unique: bool = False) -> Array_1d[np.complex128]: ... +@overload +def union1d(ar1: _ArrayLike[_AnyScalarT], ar2: _ArrayLike[_AnyScalarT]) -> Array_1d[_AnyScalarT]: ... +@overload +def union1d(ar1: ArrayLike, ar2: ArrayLike) -> Array_1d[Any]: ... + +# +@overload +def setxor1d(ar1: ToFloat64_nd, ar2: CoFloat64_nd, assume_unique: bool = False) -> Array_1d[np.float64]: ... +@overload +def setxor1d(ar1: CoFloat64_nd, ar2: ToFloat64_nd, assume_unique: bool = False) -> Array_1d[np.float64]: ... +@overload +def setxor1d(ar1: CoBool_nd, ar2: CoBool_nd, assume_unique: bool = False) -> Array_1d[np.bool]: ... +@overload +def setxor1d(ar1: ToIntP_nd, ar2: CoIntP_nd, assume_unique: bool = False) -> Array_1d[np.intp]: ... +@overload +def setxor1d(ar1: CoIntP_nd, ar2: ToIntP_nd, assume_unique: bool = False) -> Array_1d[np.intp]: ... +@overload +def setxor1d(ar1: ToComplex128_nd, ar2: CoComplex128_nd, assume_unique: bool = False) -> Array_1d[np.complex128]: ... +@overload +def setxor1d(ar1: CoComplex128_nd, ar2: ToComplex128_nd, assume_unique: bool = False) -> Array_1d[np.complex128]: ... +@overload +def setxor1d( ar1: _ArrayLike[_AnyScalarT], ar2: _ArrayLike[_AnyScalarT], - assume_unique: bool = ..., - return_indices: L[True] = ..., -) -> tuple[NDArray[_AnyScalarT], NDArray[np.int_], NDArray[np.int_]]: ... + assume_unique: bool = False, +) -> Array_1d[_AnyScalarT]: ... @overload -def intersect1d( - ar1: ArrayLike, - ar2: ArrayLike, - assume_unique: bool = ..., - return_indices: L[True] = ..., -) -> tuple[NDArray[Any], NDArray[np.int_], NDArray[np.int_]]: ... +def setxor1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: bool = False) -> Array_1d[Any]: ... # @overload -def setxor1d(ar1: _ArrayLike[_AnyScalarT], ar2: _ArrayLike[_AnyScalarT], assume_unique: bool = ...) -> NDArray[_AnyScalarT]: ... +def setdiff1d(ar1: ToFloat64_nd, ar2: CoFloat64_nd, assume_unique: bool = False) -> Array_1d[np.float64]: ... +@overload +def setdiff1d(ar1: CoFloat64_nd, ar2: ToFloat64_nd, assume_unique: bool = False) -> Array_1d[np.float64]: ... +@overload +def setdiff1d(ar1: CoBool_nd, ar2: CoBool_nd, assume_unique: bool = False) -> Array_1d[np.bool]: ... @overload -def setxor1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: bool = ...) -> NDArray[Any]: ... +def setdiff1d(ar1: ToIntP_nd, ar2: CoIntP_nd, assume_unique: bool = False) -> Array_1d[np.intp]: ... +@overload +def setdiff1d(ar1: CoIntP_nd, ar2: ToIntP_nd, assume_unique: bool = False) -> Array_1d[np.intp]: ... +@overload +def setdiff1d(ar1: ToComplex128_nd, ar2: CoComplex128_nd, assume_unique: bool = False) -> Array_1d[np.complex128]: ... +@overload +def setdiff1d(ar1: CoComplex128_nd, ar2: ToComplex128_nd, assume_unique: bool = False) -> Array_1d[np.complex128]: ... +@overload +def setdiff1d( + ar1: _ArrayLike[_AnyScalarT], + ar2: _ArrayLike[_AnyScalarT], + assume_unique: bool = False, +) -> Array_1d[_AnyScalarT]: ... +@overload +def setdiff1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: bool = False) -> Array_1d[Any]: ... # def isin( element: ArrayLike, test_elements: ArrayLike, - assume_unique: bool = ..., - invert: bool = ..., + assume_unique: bool = False, + invert: bool = False, *, - kind: str | None = ..., -) -> NDArray[np.bool]: ... + kind: L["sort", "table"] | None = None, +) -> Array[np.bool]: ... # @deprecated("Use 'isin' instead") def in1d( - element: ArrayLike, - test_elements: ArrayLike, - assume_unique: bool = ..., - invert: bool = ..., + ar1: ArrayLike, + ar2: ArrayLike, + assume_unique: bool = False, + invert: bool = False, *, - kind: str | None = ..., -) -> NDArray[np.bool]: ... - -# -@overload -def union1d(ar1: _ArrayLike[_AnyScalarT], ar2: _ArrayLike[_AnyScalarT]) -> NDArray[_AnyScalarT]: ... -@overload -def union1d(ar1: ArrayLike, ar2: ArrayLike) -> NDArray[Any]: ... -@overload -def setdiff1d(ar1: _ArrayLike[_AnyScalarT], ar2: _ArrayLike[_AnyScalarT], assume_unique: bool = ...) -> NDArray[_AnyScalarT]: ... -@overload -def setdiff1d(ar1: ArrayLike, ar2: ArrayLike, assume_unique: bool = ...) -> NDArray[Any]: ... + kind: L["sort", "table"] | None = None, +) -> Array_1d[np.bool]: ... diff --git a/test/static/accept/arraysetops.pyi b/test/static/accept/arraysetops.pyi index 44eefdb4..c4597afc 100644 --- a/test/static/accept/arraysetops.pyi +++ b/test/static/accept/arraysetops.pyi @@ -3,61 +3,63 @@ from typing_extensions import assert_type import numpy as np import numpy.typing as npt +from _numtype import Array_1d from numpy.lib._arraysetops_impl import UniqueAllResult, UniqueCountsResult, UniqueInverseResult AR_b: npt.NDArray[np.bool] -AR_i8: npt.NDArray[np.int64] +AR_i_: npt.NDArray[np.intp] AR_f8: npt.NDArray[np.float64] AR_M: npt.NDArray[np.datetime64] AR_O: npt.NDArray[np.object_] +AR_LIKE_i_: list[int] -AR_LIKE_f8: list[float] +### -assert_type(np.ediff1d(AR_b), npt.NDArray[np.int8]) -assert_type(np.ediff1d(AR_i8, to_end=[1, 2, 3]), npt.NDArray[np.int64]) -assert_type(np.ediff1d(AR_M), npt.NDArray[np.timedelta64]) -assert_type(np.ediff1d(AR_O), npt.NDArray[np.object_]) -assert_type(np.ediff1d(AR_LIKE_f8, to_begin=[1, 1.5]), npt.NDArray[Any]) - -assert_type(np.intersect1d(AR_i8, AR_i8), npt.NDArray[np.int64]) -assert_type(np.intersect1d(AR_M, AR_M, assume_unique=True), npt.NDArray[np.datetime64]) -assert_type(np.intersect1d(AR_f8, AR_i8), npt.NDArray[Any]) +assert_type(np.intersect1d(AR_i_, AR_i_), Array_1d[np.intp]) +assert_type(np.intersect1d(AR_f8, AR_i_), Array_1d[np.float64]) +assert_type(np.intersect1d(AR_M, AR_M, assume_unique=True), Array_1d[np.datetime64]) assert_type( np.intersect1d(AR_f8, AR_f8, return_indices=True), - tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]], + tuple[Array_1d[np.float64], Array_1d[np.intp], Array_1d[np.intp]], ) -assert_type(np.setxor1d(AR_i8, AR_i8), npt.NDArray[np.int64]) -assert_type(np.setxor1d(AR_M, AR_M, assume_unique=True), npt.NDArray[np.datetime64]) -assert_type(np.setxor1d(AR_f8, AR_i8), npt.NDArray[Any]) +assert_type(np.union1d(AR_i_, AR_i_), Array_1d[np.intp]) +assert_type(np.union1d(AR_f8, AR_i_), Array_1d[np.float64]) +assert_type(np.union1d(AR_M, AR_M), Array_1d[np.datetime64]) -assert_type(np.isin(AR_i8, AR_i8), npt.NDArray[np.bool]) -assert_type(np.isin(AR_M, AR_M, assume_unique=True), npt.NDArray[np.bool]) -assert_type(np.isin(AR_f8, AR_i8), npt.NDArray[np.bool]) -assert_type(np.isin(AR_f8, AR_LIKE_f8, invert=True), npt.NDArray[np.bool]) +assert_type(np.ediff1d(AR_b), Array_1d[np.int8]) +assert_type(np.ediff1d(AR_M), Array_1d[np.timedelta64]) +assert_type(np.ediff1d(AR_O), Array_1d[np.object_]) +assert_type(np.ediff1d(AR_i_, to_end=[1, 2, 3]), Array_1d[np.intp]) +assert_type(np.ediff1d(AR_LIKE_i_, to_begin=[0, 1]), Array_1d[np.intp]) -assert_type(np.union1d(AR_i8, AR_i8), npt.NDArray[np.int64]) -assert_type(np.union1d(AR_M, AR_M), npt.NDArray[np.datetime64]) -assert_type(np.union1d(AR_f8, AR_i8), npt.NDArray[Any]) +assert_type(np.setxor1d(AR_i_, AR_i_), Array_1d[np.intp]) +assert_type(np.setxor1d(AR_f8, AR_i_), Array_1d[np.float64]) +assert_type(np.setxor1d(AR_M, AR_M, assume_unique=True), Array_1d[np.datetime64]) -assert_type(np.setdiff1d(AR_i8, AR_i8), npt.NDArray[np.int64]) -assert_type(np.setdiff1d(AR_M, AR_M, assume_unique=True), npt.NDArray[np.datetime64]) -assert_type(np.setdiff1d(AR_f8, AR_i8), npt.NDArray[Any]) +assert_type(np.setdiff1d(AR_i_, AR_i_), Array_1d[np.intp]) +assert_type(np.setdiff1d(AR_f8, AR_i_), Array_1d[np.float64]) +assert_type(np.setdiff1d(AR_M, AR_M, assume_unique=True), Array_1d[np.datetime64]) + +assert_type(np.isin(AR_i_, AR_i_), npt.NDArray[np.bool]) +assert_type(np.isin(AR_f8, AR_i_), npt.NDArray[np.bool]) +assert_type(np.isin(AR_M, AR_M, assume_unique=True), npt.NDArray[np.bool]) +assert_type(np.isin(AR_f8, AR_LIKE_i_, invert=True), npt.NDArray[np.bool]) assert_type(np.unique(AR_f8), npt.NDArray[np.float64]) -assert_type(np.unique(AR_LIKE_f8, axis=0), npt.NDArray[Any]) assert_type(np.unique(AR_f8, return_index=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]]) -assert_type(np.unique(AR_LIKE_f8, return_index=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) assert_type(np.unique(AR_f8, return_inverse=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]]) -assert_type(np.unique(AR_LIKE_f8, return_inverse=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) assert_type(np.unique(AR_f8, return_counts=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp]]) -assert_type(np.unique(AR_LIKE_f8, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) +assert_type(np.unique(AR_LIKE_i_, axis=0), npt.NDArray[Any]) +assert_type(np.unique(AR_LIKE_i_, return_index=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) +assert_type(np.unique(AR_LIKE_i_, return_inverse=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) +assert_type(np.unique(AR_LIKE_i_, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp]]) assert_type( np.unique(AR_f8, return_index=True, return_inverse=True), tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]], ) assert_type( - np.unique(AR_LIKE_f8, return_index=True, return_inverse=True), + np.unique(AR_LIKE_i_, return_index=True, return_inverse=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]], ) assert_type( @@ -65,7 +67,7 @@ assert_type( tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]], ) assert_type( - np.unique(AR_LIKE_f8, return_index=True, return_counts=True), + np.unique(AR_LIKE_i_, return_index=True, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]], ) assert_type( @@ -73,7 +75,7 @@ assert_type( tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp]], ) assert_type( - np.unique(AR_LIKE_f8, return_inverse=True, return_counts=True), + np.unique(AR_LIKE_i_, return_inverse=True, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp]], ) assert_type( @@ -81,15 +83,15 @@ assert_type( tuple[npt.NDArray[np.float64], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]], ) assert_type( - np.unique(AR_LIKE_f8, return_index=True, return_inverse=True, return_counts=True), + np.unique(AR_LIKE_i_, return_index=True, return_inverse=True, return_counts=True), tuple[npt.NDArray[Any], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.intp]], ) assert_type(np.unique_all(AR_f8), UniqueAllResult[np.float64]) -assert_type(np.unique_all(AR_LIKE_f8), UniqueAllResult[Any]) assert_type(np.unique_counts(AR_f8), UniqueCountsResult[np.float64]) -assert_type(np.unique_counts(AR_LIKE_f8), UniqueCountsResult[Any]) assert_type(np.unique_inverse(AR_f8), UniqueInverseResult[np.float64]) -assert_type(np.unique_inverse(AR_LIKE_f8), UniqueInverseResult[Any]) -assert_type(np.unique_values(AR_f8), npt.NDArray[np.float64]) -assert_type(np.unique_values(AR_LIKE_f8), npt.NDArray[Any]) +assert_type(np.unique_values(AR_f8), Array_1d[np.float64]) +assert_type(np.unique_all(AR_LIKE_i_), UniqueAllResult[np.intp]) +assert_type(np.unique_counts(AR_LIKE_i_), UniqueCountsResult[np.intp]) +assert_type(np.unique_inverse(AR_LIKE_i_), UniqueInverseResult[np.intp]) +assert_type(np.unique_values(AR_LIKE_i_), Array_1d[np.intp])