TST: Add typing tests for `np.core.einsumfunc`

3 files changed, 120 insertions, 0 deletions

diff --git a/numpy/typing/tests/data/fail/einsumfunc.py b/numpy/typing/tests/data/fail/einsumfunc.py
new file mode 100644
index 000000000..33722f861
--- /dev/null
+++ b/numpy/typing/tests/data/fail/einsumfunc.py
@@ -0,0 +1,15 @@
+from typing import List, Any
+import numpy as np
+
+AR_i: np.ndarray[Any, np.dtype[np.int64]]
+AR_f: np.ndarray[Any, np.dtype[np.float64]]
+AR_m: np.ndarray[Any, np.dtype[np.timedelta64]]
+AR_O: np.ndarray[Any, np.dtype[np.object_]]
+AR_U: np.ndarray[Any, np.dtype[np.str_]]
+
+np.einsum("i,i->i", AR_i, AR_m)  # E: incompatible type
+np.einsum("i,i->i", AR_O, AR_O)  # E: incompatible type
+np.einsum("i,i->i", AR_f, AR_f, dtype=np.int32)  # E: incompatible type
+np.einsum("i,i->i", AR_i, AR_i, dtype=np.timedelta64, casting="unsafe")  # E: No overload variant
+np.einsum("i,i->i", AR_i, AR_i, out=AR_U)  # E: Value of type variable "_ArrayType" of "einsum" cannot be
+np.einsum("i,i->i", AR_i, AR_i, out=AR_U, casting="unsafe")  # E: No overload variant
diff --git a/numpy/typing/tests/data/pass/einsumfunc.py b/numpy/typing/tests/data/pass/einsumfunc.py
new file mode 100644
index 000000000..914eed4cc
--- /dev/null
+++ b/numpy/typing/tests/data/pass/einsumfunc.py
@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import List, Any
+
+import pytest
+import numpy as np
+
+AR_LIKE_b = [True, True, True]
+AR_LIKE_u = [np.uint32(1), np.uint32(2), np.uint32(3)]
+AR_LIKE_i = [1, 2, 3]
+AR_LIKE_f = [1.0, 2.0, 3.0]
+AR_LIKE_c = [1j, 2j, 3j]
+AR_LIKE_U = ["1", "2", "3"]
+
+OUT_c: np.ndarray[Any, np.dtype[np.complex128]] = np.empty(3, dtype=np.complex128)
+
+np.einsum("i,i->i", AR_LIKE_b, AR_LIKE_b)
+np.einsum("i,i->i", AR_LIKE_u, AR_LIKE_u)
+np.einsum("i,i->i", AR_LIKE_i, AR_LIKE_i)
+np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f)
+np.einsum("i,i->i", AR_LIKE_c, AR_LIKE_c)
+np.einsum("i,i->i", AR_LIKE_b, AR_LIKE_i)
+np.einsum("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c)
+
+np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f, dtype="c16")
+np.einsum("i,i->i", AR_LIKE_U, AR_LIKE_U, dtype=bool, casting="unsafe")
+np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f, out=OUT_c)
+with pytest.raises(np.ComplexWarning):
+    np.einsum("i,i->i", AR_LIKE_U, AR_LIKE_U, dtype=float, casting="unsafe", out=OUT_c)
+
+np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_b)
+np.einsum_path("i,i->i", AR_LIKE_u, AR_LIKE_u)
+np.einsum_path("i,i->i", AR_LIKE_i, AR_LIKE_i)
+np.einsum_path("i,i->i", AR_LIKE_f, AR_LIKE_f)
+np.einsum_path("i,i->i", AR_LIKE_c, AR_LIKE_c)
+np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_i)
+np.einsum_path("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c)
diff --git a/numpy/typing/tests/data/reveal/einsumfunc.py b/numpy/typing/tests/data/reveal/einsumfunc.py
new file mode 100644
index 000000000..c745a9f85
--- /dev/null
+++ b/numpy/typing/tests/data/reveal/einsumfunc.py
@@ -0,0 +1,68 @@
+from typing import List, Any
+import numpy as np
+
+AR_LIKE_b: List[bool]
+AR_LIKE_u: List[np.uint32]
+AR_LIKE_i: List[int]
+AR_LIKE_f: List[float]
+AR_LIKE_c: List[complex]
+AR_LIKE_U: List[str]
+
+OUT_f: np.ndarray[Any, np.dtype[np.float64]]
+
+reveal_type(  # E: Union[numpy.bool_, numpy.ndarray[Any, numpy.dtype[numpy.bool_]]
+    np.einsum("i,i->i", AR_LIKE_b, AR_LIKE_b)
+)
+reveal_type(  # E: Union[numpy.unsignedinteger[Any], numpy.ndarray[Any, numpy.dtype[numpy.unsignedinteger[Any]]]
+    np.einsum("i,i->i", AR_LIKE_u, AR_LIKE_u)
+)
+reveal_type(  # E: Union[numpy.signedinteger[Any], numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+    np.einsum("i,i->i", AR_LIKE_i, AR_LIKE_i)
+)
+reveal_type(  # E: Union[numpy.floating[Any], numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]]
+    np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f)
+)
+reveal_type(  # E: Union[numpy.complexfloating[Any, Any], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+    np.einsum("i,i->i", AR_LIKE_c, AR_LIKE_c)
+)
+reveal_type(  # E: Union[numpy.signedinteger[Any], numpy.ndarray[Any, numpy.dtype[numpy.signedinteger[Any]]]
+    np.einsum("i,i->i", AR_LIKE_b, AR_LIKE_i)
+)
+reveal_type(  # E: Union[numpy.complexfloating[Any, Any], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+    np.einsum("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c)
+)
+
+reveal_type(  # E: Union[{float64}, numpy.ndarray[Any, numpy.dtype[{float64}]]
+    np.einsum("i,i->i", AR_LIKE_c, AR_LIKE_c, out=OUT_f)
+)
+reveal_type(  # E: numpy.ndarray[Any, numpy.dtype[numpy.str_]
+    np.einsum("i,i->i", AR_LIKE_U, AR_LIKE_U, dtype=bool, casting="unsafe", out=OUT_f)
+)
+reveal_type(  # E: Union[numpy.complexfloating[Any, Any], numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]]
+    np.einsum("i,i->i", AR_LIKE_f, AR_LIKE_f, dtype="c16")
+)
+reveal_type(  # E: Any
+    np.einsum("i,i->i", AR_LIKE_U, AR_LIKE_U, dtype=bool, casting="unsafe")
+)
+
+reveal_type(  # E: Tuple[builtins.list[Any], builtins.str]
+    np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_b)
+)
+reveal_type(  # E: Tuple[builtins.list[Any], builtins.str]
+    np.einsum_path("i,i->i", AR_LIKE_u, AR_LIKE_u)
+)
+reveal_type(  # E: Tuple[builtins.list[Any], builtins.str]
+    np.einsum_path("i,i->i", AR_LIKE_i, AR_LIKE_i)
+)
+reveal_type(  # E: Tuple[builtins.list[Any], builtins.str]
+    np.einsum_path("i,i->i", AR_LIKE_f, AR_LIKE_f)
+)
+reveal_type(  # E: Tuple[builtins.list[Any], builtins.str]
+    np.einsum_path("i,i->i", AR_LIKE_c, AR_LIKE_c)
+)
+reveal_type(  # E: Tuple[builtins.list[Any], builtins.str]
+    np.einsum_path("i,i->i", AR_LIKE_b, AR_LIKE_i)
+)
+reveal_type(  # E: Tuple[builtins.list[Any], builtins.str]
+    np.einsum_path("i,i,i,i->i", AR_LIKE_b, AR_LIKE_u, AR_LIKE_i, AR_LIKE_c)
+)