From 59e79c3c7c6f2d1afc3db0493722f7c36fd39a84 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Tue, 9 Aug 2022 16:08:05 +0300
Subject: [PATCH] PY-54503 Provide the result type for Enum[...].value and
 Enum(...).value

Precise types can be inferred only over AST if it's accessible.

I had to move PyStdlibTypeProvider higher in the provider's hierarchy
so that it could override types coming from Typeshed, otherwise we
infer enum.property type for the "value" attribute.

GitOrigin-RevId: 8727e080cfc06d0edda13eccfd601601dc661da9
---
 .../resources/META-INF/PythonPsiImpl.xml      |  3 +-
 .../stdlib/PyStdlibTypeProvider.java          | 21 ++++++++++---
 .../com/jetbrains/python/PyTypeTest.java      | 30 +++++++++++++++++++
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml b/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml
index c60c79523a82..fff4307cdbf2 100644
--- a/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml
+++ b/python/python-psi-impl/resources/META-INF/PythonPsiImpl.xml
@@ -523,7 +523,8 @@
     <inspectionExtension implementation="com.jetbrains.python.inspections.stdlib.PyStdlibInspectionExtension"/>
     <typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyDataclassTypeProvider"/>
     <typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyNamedTupleTypeProvider"/>
-    <typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibTypeProvider"/>
+    <!-- This provider should be able to override the results on anything else, including types coming from .pyi stubs -->
+    <typeProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibTypeProvider" order="first"/>
     <pyClassMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyOverridingClassDunderMembersProvider"/>
     <pyClassMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibClassMembersProvider"/>
     <pyModuleMembersProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibOverridingModuleMembersProvider"/>
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 25437ef8fe99..13860f1d7b30 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
@@ -110,10 +110,23 @@ public class PyStdlibTypeProvider extends PyTypeProviderBase {
           final PyReferenceExpression qualifierExpr = (PyReferenceExpression)qualifier;
           final PsiElement resolvedQualifier = qualifierExpr.getReference().resolve();
           if (resolvedQualifier instanceof PyTargetExpression) {
-            final PyTargetExpression qualifierTarget = (PyTargetExpression)resolvedQualifier;
-            // Requires switching to AST, we cannot use getType(qualifierTarget) here, because its type is overridden by this type provider
-            if (context.maySwitchToAST(qualifierTarget)) {
-              final PyExpression value = qualifierTarget.findAssignedValue();
+            final PyTargetExpression enumItem = (PyTargetExpression)resolvedQualifier;
+            // Requires switching to AST, we cannot use getType(enumItem) here, because its type is overridden by this type provider
+            if (context.maySwitchToAST(enumItem)) {
+              final PyExpression value = enumItem.findAssignedValue();
+              if (value != null) {
+                return context.getType(value);
+              }
+            }
+          }
+        }
+        else if (qualifier != null) {
+          PyClassType enumType = as(context.getType(qualifier), PyClassType.class);
+          if (enumType != null) {
+            PyClass enumClass = enumType.getPyClass();
+            PyTargetExpression firstEnumItem = ContainerUtil.getFirstItem(enumClass.getClassAttributes());
+            if (firstEnumItem != null && context.maySwitchToAST(firstEnumItem)) {
+              final PyExpression value = firstEnumItem.findAssignedValue();
               if (value != null) {
                 return context.getType(value);
               }
diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java
index 8427150d31eb..251020dcfdaf 100644
--- a/python/testSrc/com/jetbrains/python/PyTypeTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java
@@ -3979,6 +3979,36 @@ public class PyTypeTest extends PyTestCase {
     );
   }
 
+  // PY-54503
+  public void testEnumGetItemResultValueAttribute() {
+    runWithLanguageLevel(
+      LanguageLevel.getLatest(),
+      () -> doTest("int",
+                   "import enum\n" +
+                   "\n" +
+                   "class MyEnum(enum.Enum):\n" +
+                   "    ONE = 1\n" +
+                   "    TWO = 2\n" +
+                   "\n" +
+                   "expr = MyEnum['ONE'].value")
+    );
+  }
+
+  // PY-54503
+  public void testEnumDunderCallResultValueAttribute() {
+    runWithLanguageLevel(
+      LanguageLevel.getLatest(),
+      () -> doTest("int",
+                   "import enum\n" +
+                   "\n" +
+                   "class MyEnum(enum.Enum):\n" +
+                   "    ONE = 1\n" +
+                   "    TWO = 2\n" +
+                   "\n" +
+                   "expr = MyEnum(1).value")
+    );
+  }
+
   private static List<TypeEvalContext> getTypeEvalContexts(@NotNull PyExpression element) {
     return ImmutableList.of(TypeEvalContext.codeAnalysis(element.getProject(), element.getContainingFile()).withTracing(),
                             TypeEvalContext.userInitiated(element.getProject(), element.getContainingFile()).withTracing());