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(
"""