🏷️ Annotate einsum_path return type (Issue #724) (#746)

Author: pareshjoshij

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])

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]: ...