From cc99207f78c84a7c7a89c5c1523371d1dcfe23bf Mon Sep 17 00:00:00 2001
From: Bas van Beek <bas.vanbeek@hotmail.com>
Date: Mon, 8 May 2023 15:31:33 +0200
Subject: TYP: Let `np.einsum` accept `object` dtypes

xref https://github.com/numpy/numpy/pull/18053
---
 numpy/core/einsumfunc.pyi | 53 ++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 48 insertions(+), 5 deletions(-)

(limited to 'numpy/core')

diff --git a/numpy/core/einsumfunc.pyi b/numpy/core/einsumfunc.pyi
index c811a5783..ad483bb90 100644
--- a/numpy/core/einsumfunc.pyi
+++ b/numpy/core/einsumfunc.pyi
@@ -5,10 +5,6 @@ from numpy import (
     ndarray,
     dtype,
     bool_,
-    unsignedinteger,
-    signedinteger,
-    floating,
-    complexfloating,
     number,
     _OrderKACF,
 )
@@ -18,12 +14,14 @@ from numpy._typing import (
     _ArrayLikeInt_co,
     _ArrayLikeFloat_co,
     _ArrayLikeComplex_co,
+    _ArrayLikeObject_co,
     _DTypeLikeBool,
     _DTypeLikeUInt,
     _DTypeLikeInt,
     _DTypeLikeFloat,
     _DTypeLikeComplex,
     _DTypeLikeComplex_co,
+    _DTypeLikeObject,
 )
 
 _ArrayType = TypeVar(
@@ -132,6 +130,51 @@ def einsum(
     optimize: _OptimizeKind = ...,
 ) -> _ArrayType: ...
 
+@overload
+def einsum(
+    subscripts: str | _ArrayLikeInt_co,
+    /,
+    *operands: _ArrayLikeObject_co,
+    out: None = ...,
+    dtype: None | _DTypeLikeObject = ...,
+    order: _OrderKACF = ...,
+    casting: _CastingSafe = ...,
+    optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+    subscripts: str | _ArrayLikeInt_co,
+    /,
+    *operands: Any,
+    casting: _CastingUnsafe,
+    dtype: None | _DTypeLikeObject = ...,
+    out: None = ...,
+    order: _OrderKACF = ...,
+    optimize: _OptimizeKind = ...,
+) -> Any: ...
+@overload
+def einsum(
+    subscripts: str | _ArrayLikeInt_co,
+    /,
+    *operands: _ArrayLikeObject_co,
+    out: _ArrayType,
+    dtype: None | _DTypeLikeObject = ...,
+    order: _OrderKACF = ...,
+    casting: _CastingSafe = ...,
+    optimize: _OptimizeKind = ...,
+) -> _ArrayType: ...
+@overload
+def einsum(
+    subscripts: str | _ArrayLikeInt_co,
+    /,
+    *operands: Any,
+    out: _ArrayType,
+    casting: _CastingUnsafe,
+    dtype: None | _DTypeLikeObject = ...,
+    order: _OrderKACF = ...,
+    optimize: _OptimizeKind = ...,
+) -> _ArrayType: ...
+
 # NOTE: `einsum_call` is a hidden kwarg unavailable for public use.
 # It is therefore excluded from the signatures below.
 # NOTE: In practice the list consists of a `str` (first element)
@@ -139,6 +182,6 @@ def einsum(
 def einsum_path(
     subscripts: str | _ArrayLikeInt_co,
     /,
-    *operands: _ArrayLikeComplex_co,
+    *operands: _ArrayLikeComplex_co | _DTypeLikeObject,
     optimize: _OptimizeKind = ...,
 ) -> tuple[list[Any], str]: ...
-- 
cgit v1.2.1