From 5579292ff37c4c574e812164922c7c8cd3b1d88f Mon Sep 17 00:00:00 2001 From: chbndrhnns Date: Thu, 30 Oct 2025 21:30:33 +0100 Subject: [PATCH] [python] PY-59260 False positive assigining `IntFlag` enum value to an annotated variable (cherry picked from commit 1b231b0917c41032255e90b3f524fc3d6201e80e) IJ-CR-181224 GitOrigin-RevId: e2969be0fd18d29a1c2ba00a7ac3ad1a87edf79e --- .../stdlib/PyStdlibTypeProvider.java | 58 +++++++-- .../PyTypeCheckerInspectionTest.java | 112 ++++++++++++++++++ 2 files changed, 160 insertions(+), 10 deletions(-) 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 ecb7a763300b..1511dfbb3ec8 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 @@ -109,6 +109,12 @@ public final class PyStdlibTypeProvider extends PyTypeProviderBase { else if ("enum.IntEnum.value".equals(name) && anchor instanceof PyReferenceExpression) { return Ref.create(PyBuiltinCache.getInstance(referenceTarget).getIntType()); } + else if ("enum.StrEnum.value".equals(name) && anchor instanceof PyReferenceExpression) { + return Ref.create(PyBuiltinCache.getInstance(referenceTarget).getStrType()); + } + else if ((PyNames.TYPE_ENUM_FLAG + ".value").equals(name)&& anchor instanceof PyReferenceExpression) { + return Ref.create(PyBuiltinCache.getInstance(referenceTarget).getIntType()); + } else if ((PyNames.TYPE_ENUM + ".value").equals(name) && anchor instanceof PyReferenceExpression anchorExpr && context.maySwitchToAST(anchor)) { final PyExpression qualifier = anchorExpr.getQualifier(); @@ -118,16 +124,7 @@ public final class PyStdlibTypeProvider extends PyTypeProviderBase { if (enumType != null) { PyClass enumClass = enumType.getPyClass(); if (isCustomEnum(enumClass, context)) { - PyTargetExpression firstEnumItem = ContainerUtil.getFirstItem(enumClass.getClassAttributes()); - if (firstEnumItem != null) { - EnumAttributeInfo attributeInfo = getEnumAttributeInfo(enumClass, firstEnumItem, context); - if (attributeInfo != null) { - return Ref.create(attributeInfo.assignedValueType); - } - } - else { - return Ref.create(); - } + return Ref.create(getEnumValueType(enumClass, context)); } } } @@ -348,6 +345,47 @@ public final class PyStdlibTypeProvider extends PyTypeProviderBase { return null; } + // Handle IntEnum/IntFlag, StrEnum, and fall back to assigned type or unknown + private static @Nullable PyType getEnumValueType(@NotNull PyClass enumClass, @NotNull TypeEvalContext context) { + PyBuiltinCache cache = PyBuiltinCache.getInstance(enumClass); + + if (enumClass.isSubclass("enum.IntEnum", context) || + enumClass.isSubclass("enum.IntFlag", context) || + enumClass.isSubclass("enum.Flag", context)) { + return cache.getIntType(); + } + if (enumClass.isSubclass("enum.StrEnum", context)) { + return cache.getStrType(); + } + + // If it's a mixin with basic scalar types - check str first for the failing test + if (enumClass.isSubclass("builtins.str", context) || enumClass.isSubclass("str", context)) { + return cache.getStrType(); + } + if (enumClass.isSubclass("builtins.int", context) || enumClass.isSubclass("int", context)) { + return cache.getIntType(); + } + if (enumClass.isSubclass("builtins.bytes", context) || enumClass.isSubclass("bytes", context)) { + return cache.getBytesType(LanguageLevel.forElement(enumClass)); + } + if (enumClass.isSubclass("builtins.float", context) || enumClass.isSubclass("float", context)) { + return cache.getFloatType(); + } + if (enumClass.isSubclass("builtins.bool", context) || enumClass.isSubclass("bool", context)) { + return cache.getBoolType(); + } + + // Fallback: Infer from first MEMBER's assigned value (not non-members like helpers/descriptors) + for (PyTargetExpression targetExpr : enumClass.getClassAttributes()) { + EnumAttributeInfo attributeInfo = getEnumAttributeInfo(enumClass, targetExpr, context); + if (attributeInfo != null && attributeInfo.attributeKind == EnumAttributeKind.MEMBER) { + return attributeInfo.assignedValueType; + } + } + + return null; + } + @Override public @Nullable Ref getCallType(@NotNull PyFunction function, @NotNull PyCallSiteExpression callSite, @NotNull TypeEvalContext context) { final String qname = function.getQualifiedName(); diff --git a/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java index 8cce3b44e7e3..9d536c6a7dff 100644 --- a/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java +++ b/python/testSrc/com/jetbrains/python/inspections/PyTypeCheckerInspectionTest.java @@ -1618,4 +1618,116 @@ public class PyTypeCheckerInspectionTest extends PyInspectionTestCase { z: type[None] = NoneType """)); } + + // PY-59260 + public void testIntFlagValueType() { + runWithLanguageLevel( + LanguageLevel.PYTHON36, + () -> + doTestByText(""" + from enum import IntFlag, auto + + # IntFlag should infer int + class IF(IntFlag): + FIRST = auto() + SECOND = auto() + THIRD = 42 + + # IntFlag.value should return int, so these should not produce type errors + variable: int = IF.FIRST.value + another_var: int = IF.SECOND.value + explicit_var: int = IF.THIRD.value + + # This should produce a type error + wrong_var: str = IF.FIRST.value + """)); + } + + // PY-59260 + public void testEnumValueTypeInference() { + runWithLanguageLevel( + LanguageLevel.getLatest(), + () -> + doTestByText(""" + from enum import Enum, IntFlag, StrEnum + + # IntFlag should infer int + class IF(IntFlag): + A = 1 + i: int = IF.A.value + + # StrEnum should infer str + class SE(StrEnum): + B = "b" + s: str = SE.B.value + + # str mixin should infer str + class StrMixin(str, Enum): + C = "c" + s2: str = StrMixin.C.value + s3: int = StrMixin.C.value + + # Empty str mixin should also infer str + class EmptyStrMixin(str, Enum): + pass + def test_empty(x: EmptyStrMixin): + s4: str = x.value + i2: int = x.value + """)); + } + + // PY-59260 + public void testEmptyEnumValueTypes() { + runWithLanguageLevel( + LanguageLevel.PYTHON311, + () -> + doTestByText(""" + from enum import StrEnum, Enum + + class EmptyStrEnum(StrEnum): + pass + + class EmptyStrMixin(str, Enum): + pass + + def test_empty_str_enum(x: EmptyStrEnum): + s: str = x.value + i: int = x.value + + def test_empty_str_mixin(x: EmptyStrMixin): + s: str = x.value + i: int = x.value + """)); + } + + // PY-59260 + public void testEnumValueTypeIgnoresNonMembers() { + runWithLanguageLevel( + LanguageLevel.getLatest(), + () -> + doTestByText(""" + from enum import Enum, nonmember + + # Simple enum with just integer members + class SimpleEnum(Enum): + A = 1 + B = 2 + + # Should infer int + x: int = SimpleEnum.A.value + y: str = SimpleEnum.B.value + + # Enum with non-member first, then actual members + class E(Enum): + # This should be classified as a non-member + HELPER_CONSTANT = nonmember("not a member") + # These are the actual members - should infer int from first member + FIRST_MEMBER = 42 + SECOND_MEMBER = 43 + + # Should infer int from FIRST_MEMBER (ignoring HELPER_CONSTANT) + a: int = E.FIRST_MEMBER.value + b: str = E.SECOND_MEMBER.value + """)); + } }