PY-80565 False warning auto not assignable to StrEnum

GitOrigin-RevId: 6e9bf2dad1ca9d421d433def0d1c323a1fd3b301

python/python-psi-impl/src/com/jetbrains/python/codeInsight/stdlib/PyStdlibTypeProvider.java

@@ -249,48 +249,66 @@ public final class PyStdlibTypeProvider extends PyTypeProviderBase {
     }
   }
 
-  private static @NotNull EnumAttributeInfo getEnumAttributeInfo(@NotNull PyClass enumClass, @Nullable PyType type, @NotNull TypeEvalContext context) {
-    if (type == null) {
-      return new EnumAttributeInfo(null, EnumAttributeKind.MEMBER);
+  private static @NotNull EnumAttributeInfo getEnumAttributeInfo(@NotNull PyClass enumClass,
+                                                                 @Nullable PyType type,
+                                                                 @NotNull TypeEvalContext context) {
+    Boolean isMember = null;
+    if (type instanceof PyCollectionType genericType) {
+      if (PyNames.TYPE_ENUM_MEMBER.equals(genericType.getClassQName())) {
+        type = ContainerUtil.getOnlyItem(genericType.getElementTypes());
+        isMember = true;
+      }
+      if (PyNames.TYPE_ENUM_NONMEMBER.equals(genericType.getClassQName())) {
+        type = ContainerUtil.getOnlyItem(genericType.getElementTypes());
+        isMember = false;
+      }
     }
-    PyQualifiedNameOwner typeDeclarationElement = type.getDeclarationElement();
-    if (typeDeclarationElement != null) {
-      String typeDeclarationQName = typeDeclarationElement.getQualifiedName();
-      if (type instanceof PyCollectionType genericType) {
-        PyType genericParameterType = ContainerUtil.getOnlyItem(genericType.getElementTypes());
-        if (genericParameterType != null) {
-          if (PyNames.TYPE_ENUM_MEMBER.equals(typeDeclarationQName)) {
-            return EnumAttributeInfo.memberOrAlias(enumClass, genericParameterType);
-          }
-          if (PyNames.TYPE_ENUM_NONMEMBER.equals(typeDeclarationQName)) {
-            return EnumAttributeInfo.nonMember(genericParameterType);
+    if (type instanceof PyClassType classType && PyNames.TYPE_ENUM_AUTO.equals(classType.getClassQName())) {
+      type = getEnumAutoValueType(enumClass, context);
+    }
+
+    if (isMember == null) {
+      isMember = true;
+      if (type != null) {
+        if (type.getDeclarationElement() instanceof PyCallable) {
+          isMember = false;
+        }
+        else {
+          boolean isDescriptor = !ContainerUtil.isEmpty(
+            type.resolveMember(PyNames.DUNDER_GET, null, AccessDirection.READ, PyResolveContext.defaultContext(context)));
+          if (isDescriptor) {
+            isMember = false;
           }
         }
       }
     }
-    if (typeDeclarationElement instanceof PyCallable) {
-      return EnumAttributeInfo.nonMember(type);
+
+    EnumAttributeKind attributeKind;
+    if (isMember) {
+      attributeKind = type instanceof PyLiteralType literalType && literalType.getPyClass().equals(enumClass)
+                                        ? EnumAttributeKind.MEMBER_ALIAS
+                                        : EnumAttributeKind.MEMBER;
     }
-    boolean isDescriptor = !ContainerUtil.isEmpty(
-      type.resolveMember(PyNames.DUNDER_GET, null, AccessDirection.READ, PyResolveContext.defaultContext(context)));
-    if (isDescriptor) {
-      return EnumAttributeInfo.nonMember(type);
+    else {
+      attributeKind = EnumAttributeKind.NONMEMBER;
     }
-    return EnumAttributeInfo.memberOrAlias(enumClass, type);
+    return new EnumAttributeInfo(type, attributeKind);
+  }
+
+  private static @Nullable PyType getEnumAutoValueType(@NotNull PyClass enumClass, @NotNull TypeEvalContext context) {
+    PyFunction generateNextValueMethod = enumClass.findMethodByName("_generate_next_value_", true, context);
+    if (generateNextValueMethod != null) {
+      // Ignore 'Enum._generate_next_value_' as its declared return type is 'Any' in the typeshed stubs.
+      PyClass containingClass = generateNextValueMethod.getContainingClass();
+      if (containingClass != null && !PyNames.TYPE_ENUM.equals(containingClass.getQualifiedName())) {
+        return context.getReturnType(generateNextValueMethod);
+      }
+    }
+    return PyBuiltinCache.getInstance(enumClass).getIntType();
   }
 
   @ApiStatus.Internal
   public record EnumAttributeInfo(@Nullable PyType assignedValueType, @NotNull EnumAttributeKind attributeKind) {
-    private static @NotNull EnumAttributeInfo memberOrAlias(@NotNull PyClass enumClass, @Nullable PyType type) {
-      EnumAttributeKind attributeKind = type instanceof PyLiteralType literalType && literalType.getPyClass().equals(enumClass)
-                                        ? EnumAttributeKind.MEMBER_ALIAS
-                                        : EnumAttributeKind.MEMBER;
-      return new EnumAttributeInfo(type, attributeKind);
-    }
-
-    private static @NotNull EnumAttributeInfo nonMember(@Nullable PyType type) {
-      return new EnumAttributeInfo(type, EnumAttributeKind.NONMEMBER);
-    }
   }
 
   @ApiStatus.Internal

python/testData/inspections/PyEnumInspection/CustomGenerateNextValueMultiFile/a.py

@@ -0,0 +1,7 @@
+from enum import auto, Enum
+from my_enum_base import MyEnumBase
+
+
+class MyEnumDerived(MyEnumBase):
+    FOO = auto()
+    BAR = <warning descr="Type 'int' is not assignable to declared type 'MyVal'">2</warning>

python/testData/inspections/PyEnumInspection/CustomGenerateNextValueMultiFile/my_enum_base.py

@@ -0,0 +1,11 @@
+from enum import Enum
+
+
+class MyVal: ...
+
+
+class MyEnumBase(Enum):
+    _value_: MyVal
+
+    @staticmethod
+    def _generate_next_value_(name: str, start: int, count: int, last_values: list[MyVal]) -> MyVal: ...

python/testSrc/com/jetbrains/python/inspections/PyEnumInspectionTest.kt

@@ -51,19 +51,75 @@ class PyEnumInspectionTest : PyInspectionTestCase() {
   fun testDeclaredEnumMemberType() {
     doTestByText(
       """
-        from enum import Enum, auto
+        from enum import Enum
         
         class Color(Enum):
           _value_: int
           RED = 1
           GREEN = <warning descr="Type 'str' is not assignable to declared type 'int'">"green"</warning>
-          BLUE = auto()
           
           R = RED
       """.trimIndent()
     )
   }
 
+  // PY-80565
+  fun testEnumMemberAssignedToAuto() {
+    doTestByText(
+      """
+        from enum import auto, Enum
+        
+        class MyIntEnum(Enum):
+            _value_: int
+            FOO = auto()
+        
+        class MyStrEnum(Enum):
+            _value_: str
+            FOO = <warning descr="Type 'int' is not assignable to declared type 'str'">auto()</warning>
+        """.trimIndent()
+    )
+  }
+
+  // PY-80565
+  fun testCustomGenerateNextValue() {
+    doTestByText(
+      """
+        from enum import auto, Enum
+        
+        class MyVal: ...
+
+        class MyEnumBase(Enum):
+            _value_: MyVal        
+
+            @staticmethod
+            def _generate_next_value_(name: str, start: int, count: int, last_values: list[MyVal]) -> MyVal: ...
+        
+        
+        class MyEnumDerived(MyEnumBase):
+            FOO = auto()
+            BAR = <warning descr="Type 'int' is not assignable to declared type 'MyVal'">2</warning>
+      """.trimIndent()
+    )
+  }
+
+  // PY-80565
+  fun testCustomGenerateNextValueMultiFile() {
+    doMultiFileTest()
+  }
+
+  // PY-80565
+  fun testStrEnumMemberAssignedToAuto() {
+    doTestByText(
+      """
+        from enum import auto, StrEnum
+        
+        class Example(StrEnum):
+            FOO = auto()
+            BAR = <warning descr="Type 'int' is not assignable to declared type 'str'">1</warning>
+        """.trimIndent()
+    )
+  }
+
   fun testTypeAnnotationsForEnumMember() {
     doTestByText(
       """