diff --git a/src/numpy-stubs/@test/static/accept/einsumfunc.pyi b/src/numpy-stubs/@test/static/accept/einsumfunc.pyi index dc488f31..b2e57356 100644 --- a/src/numpy-stubs/@test/static/accept/einsumfunc.pyi +++ b/src/numpy-stubs/@test/static/accept/einsumfunc.pyi @@ -27,13 +27,15 @@ assert_type(np.einsum("i,i->i", AR_LIKE_U, AR_LIKE_U, dtype=bool, casting="unsaf assert_type(np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f, dtype="c16"), Any) assert_type(np.einsum("i,i->i", AR_LIKE_U, AR_LIKE_U, dtype=bool, casting="unsafe"), Any) -assert_type(np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_b), tuple[list[Any], str]) -assert_type(np.einsum_path("i,i->i", AR_LIKE_u, AR_LIKE_u), tuple[list[Any], str]) -assert_type(np.einsum_path("i,i->i", AR_LIKE_i, AR_LIKE_i), tuple[list[Any], str]) -assert_type(np.einsum_path("i,i->i", AR_LIKE_f, AR_LIKE_f), tuple[list[Any], str]) -assert_type(np.einsum_path("i,i->i", AR_LIKE_c, AR_LIKE_c), tuple[list[Any], str]) -assert_type(np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_i), tuple[list[Any], str]) -assert_type(np.einsum_path("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c), tuple[list[Any], str]) +assert_type(np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_b), tuple[list[str | tuple[int, ...]], str]) +assert_type(np.einsum_path("i,i->i", AR_LIKE_u, AR_LIKE_u), tuple[list[str | tuple[int, ...]], str]) +assert_type(np.einsum_path("i,i->i", AR_LIKE_i, AR_LIKE_i), tuple[list[str | tuple[int, ...]], str]) +assert_type(np.einsum_path("i,i->i", AR_LIKE_f, AR_LIKE_f), tuple[list[str | tuple[int, ...]], str]) +assert_type(np.einsum_path("i,i->i", AR_LIKE_c, AR_LIKE_c), tuple[list[str | tuple[int, ...]], str]) +assert_type(np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_i), tuple[list[str | tuple[int, ...]], str]) +assert_type( + np.einsum_path("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c), tuple[list[str | tuple[int, ...]], str] +) assert_type(np.einsum([[1, 1], [1, 1]], AR_LIKE_i, AR_LIKE_i), Any) -assert_type(np.einsum_path([[1, 1], [1, 1]], AR_LIKE_i, AR_LIKE_i), tuple[list[Any], str]) +assert_type(np.einsum_path([[1, 1], [1, 1]], AR_LIKE_i, AR_LIKE_i), tuple[list[str | tuple[int, ...]], str]) diff --git a/src/numpy-stubs/_core/einsumfunc.pyi b/src/numpy-stubs/_core/einsumfunc.pyi index d1ee7aeb..158d746f 100644 --- a/src/numpy-stubs/_core/einsumfunc.pyi +++ b/src/numpy-stubs/_core/einsumfunc.pyi @@ -26,7 +26,7 @@ __all__ = ["einsum", "einsum_path"] _ArrayT = TypeVar("_ArrayT", bound=_nt.Array[_nt.co_complex]) # TODO (@jorenham): Annotate the `Sequence` value (numpy/numtype#724) -_OptimizeKind: TypeAlias = bool | Literal["greedy", "optimal"] | Sequence[Incomplete] +_OptimizeKind: TypeAlias = bool | Literal["greedy", "optimal"] | Sequence[str | tuple[int, ...]] _CastingSafe: TypeAlias = Literal["no", "equiv", "safe", "same_kind"] _CastingUnsafe: TypeAlias = Literal["unsafe"] @@ -178,4 +178,4 @@ def einsum_path( *operands: _ArrayLikeComplex_co | _DTypeLikeObject, optimize: _OptimizeKind = "greedy", einsum_call: L[False] = False, -) -> tuple[list[Incomplete], str]: ... +) -> tuple[list[str | tuple[int, ...]], str]: ...