From 0fc1981b5c9ff97e018a9bf8086e21cf78a444f2 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Tue, 9 Aug 2022 12:30:51 +0300
Subject: [PATCH] PY-53388 Override the callable type for enum.auto() from
 Typeshed stubs

GitOrigin-RevId: 132e934de4da01983bebc55fcb1d3440d7053249
---
 .../stdlib/PyStdlibTypeProvider.java           | 18 ++++++++++++++++++
 .../Py3ArgumentListInspectionTest.java         | 10 ++++++++++
 2 files changed, 28 insertions(+)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java
index 61d6c5e13a9f..25437ef8fe99 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java
@@ -16,6 +16,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 
@@ -124,6 +125,23 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase {
         return PyTypeParser.getTypeByName(referenceTarget, "dict[str, unknown]", context);
       }
     }
+    @Nullable PyType enumAutoType = getEnumAutoConstructorType(referenceTarget, context, anchor);
+    if (enumAutoType != null) {
+      return enumAutoType;
+    }
+    return null;
+  }
+
+  @Nullable
+  private static PyType getEnumAutoConstructorType(@NotNull PsiElement target,
+                                                   @NotNull TypeEvalContext context,
+                                                   @Nullable PsiElement anchor) {
+    if (target instanceof PyClass && "enum.auto".equals(((PyClass)target).getQualifiedName()) && anchor instanceof PyCallExpression) {
+      PyClassLikeType classType = as(context.getType((PyTypedElement)target), PyClassLikeType.class);
+      if (classType != null) {
+        return new PyCallableTypeImpl(Collections.emptyList(), classType.toInstance());
+      }
+    }
     return null;
   }
 
diff --git a/python/testSrc/com/jetbrains/python/inspections/Py3ArgumentListInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/Py3ArgumentListInspectionTest.java
index 68be755d346d..370f7a4ba4de 100644
--- a/python/testSrc/com/jetbrains/python/inspections/Py3ArgumentListInspectionTest.java
+++ b/python/testSrc/com/jetbrains/python/inspections/Py3ArgumentListInspectionTest.java
@@ -74,4 +74,14 @@ public class Py3ArgumentListInspectionTest extends PyInspectionTestCase {
                  "\n" +
                  "random.randint(1, 2)");
   }
+
+  // PY-53388
+  public void testEnumAuto() {
+    doTestByText("import enum\n" +
+                 "\n" +
+                 "class MyEnum(enum.Enum):\n" +
+                 "    FOO = enum.auto()\n" +
+                 "    BAR = enum.auto()\n" +
+                 "");
+  }
 }