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 dd9beb38c4d2..4e873f80266e 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 @@ -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 diff --git a/python/testData/inspections/PyEnumInspection/CustomGenerateNextValueMultiFile/a.py b/python/testData/inspections/PyEnumInspection/CustomGenerateNextValueMultiFile/a.py new file mode 100644 index 000000000000..7042879fbe2f --- /dev/null +++ b/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 = 2 diff --git a/python/testData/inspections/PyEnumInspection/CustomGenerateNextValueMultiFile/my_enum_base.py b/python/testData/inspections/PyEnumInspection/CustomGenerateNextValueMultiFile/my_enum_base.py new file mode 100644 index 000000000000..9a3b0a42fcaa --- /dev/null +++ b/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: ... diff --git a/python/testSrc/com/jetbrains/python/inspections/PyEnumInspectionTest.kt b/python/testSrc/com/jetbrains/python/inspections/PyEnumInspectionTest.kt index b87fd4fbbb33..506c571e2727 100644 --- a/python/testSrc/com/jetbrains/python/inspections/PyEnumInspectionTest.kt +++ b/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 = "green" - 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 = auto() + """.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 = 2 + """.trimIndent() + ) + } + + // PY-80565 + fun testCustomGenerateNextValueMultiFile() { + doMultiFileTest() + } + + // PY-80565 + fun testStrEnumMemberAssignedToAuto() { + doTestByText( + """ + from enum import auto, StrEnum + + class Example(StrEnum): + FOO = auto() + BAR = 1 + """.trimIndent() + ) + } + fun testTypeAnnotationsForEnumMember() { doTestByText( """