diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts index 82f2353e0f72..5fb954075128 100644 --- a/packages/pyright-internal/src/analyzer/typeGuards.ts +++ b/packages/pyright-internal/src/analyzer/typeGuards.ts @@ -72,6 +72,7 @@ import { convertToInstance, convertToInstantiable, derivesFromAnyOrUnknown, + derivesFromStdlibClass, doForEachSubtype, getSpecializedTupleType, getTypeCondition, @@ -2690,6 +2691,65 @@ function isFilterSuperclass( return false; } +// Determines whether the specified class is an enum class whose members derive +// from a primitive type (e.g. IntEnum or StrEnum). Members of such classes +// compare equal to their underlying primitive value at runtime. +function isPrimitiveBackedEnumClass(classType: ClassType): boolean { + if (!ClassType.isEnumClass(classType)) { + return false; + } + + return derivesFromStdlibClass(classType, 'int') || derivesFromStdlibClass(classType, 'str'); +} + +// Determines whether enumType is a primitive-backed enum class and primitiveType +// is the primitive type (int or str) from which it derives. +function isPrimitiveBackedEnumAndPrimitive(enumType: ClassType, primitiveType: ClassType): boolean { + if (ClassType.isEnumClass(primitiveType) || !isPrimitiveBackedEnumClass(enumType)) { + return false; + } + + // Exclude bool, which derives from int but whose literal values are distinct + // from the corresponding int literals. + if (ClassType.isBuiltIn(primitiveType, 'int')) { + return derivesFromStdlibClass(enumType, 'int'); + } + + if (ClassType.isBuiltIn(primitiveType, 'str')) { + return derivesFromStdlibClass(enumType, 'str'); + } + + return false; +} + +// If the specified type is a literal member of a primitive-backed enum class, +// returns the corresponding primitive literal type. Otherwise returns the type +// unmodified. +function unwrapPrimitiveBackedEnumLiteral(classType: ClassType): ClassType { + const literalValue = classType.priv.literalValue; + if (!(literalValue instanceof EnumLiteral) || !isPrimitiveBackedEnumClass(classType)) { + return classType; + } + + const itemType = literalValue.itemType; + if (isClassInstance(itemType) && itemType.priv.literalValue !== undefined) { + return itemType; + } + + return classType; +} + +// Determines whether two literal types compare equal using the `==` operator at +// runtime. This differs from ClassType.isLiteralValueSame only for members of +// primitive-backed enum classes, which compare equal to their underlying +// primitive values. +function isLiteralValueEqualAtRuntime(type1: ClassType, type2: ClassType): boolean { + return ClassType.isLiteralValueSame( + unwrapPrimitiveBackedEnumLiteral(type1), + unwrapPrimitiveBackedEnumLiteral(type2) + ); +} + // Attempts to narrow a type (make it more constrained) based on a comparison // (equal or not equal) to a literal value. It also handles "is" or "is not" // operators if isIsOperator is true. @@ -2711,9 +2771,24 @@ function narrowTypeForLiteralComparison( return subtype; } - if (isClassInstance(subtype) && ClassType.isSameGenericClass(literalType, subtype)) { + // Determine whether this is a comparison between a primitive-backed enum + // (e.g. IntEnum or StrEnum) and its underlying primitive type. Such a + // comparison can be used for narrowing because the enum member compares + // equal to its underlying primitive value at runtime. + const isPrimitiveBackedEnumComparison = + !isIsOperator && + isClassInstance(subtype) && + (isPrimitiveBackedEnumAndPrimitive(subtype, literalType) || + isPrimitiveBackedEnumAndPrimitive(literalType, subtype)); + + if ( + isClassInstance(subtype) && + (ClassType.isSameGenericClass(literalType, subtype) || isPrimitiveBackedEnumComparison) + ) { if (subtype.priv.literalValue !== undefined) { - const literalValueMatches = ClassType.isLiteralValueSame(subtype, literalType); + const literalValueMatches = isPrimitiveBackedEnumComparison + ? isLiteralValueEqualAtRuntime(subtype, literalType) + : ClassType.isLiteralValueSame(subtype, literalType); if (isPositiveTest) { return literalValueMatches ? subtype : undefined; } @@ -2729,6 +2804,24 @@ function narrowTypeForLiteralComparison( } if (isPositiveTest) { + if (isPrimitiveBackedEnumComparison) { + // If the reference type is the enum, attempt to find the member + // whose value matches the primitive literal. If no member can be + // identified (e.g. for an IntFlag class, whose members cannot be + // enumerated), retain the reference type rather than narrowing it + // to the primitive literal type. + if (ClassType.isEnumClass(subtype)) { + const allLiteralTypes = enumerateLiteralsForType(evaluator, subtype); + const match = allLiteralTypes?.find((type) => isLiteralValueEqualAtRuntime(type, literalType)); + return match ?? subtype; + } + + // The reference type is the primitive type, so narrowing it to the + // enum literal would lose its (potentially derived) class. Retain + // the reference type instead. + return subtype; + } + return literalType; } @@ -2736,7 +2829,13 @@ function narrowTypeForLiteralComparison( // (for bool or enum), we can eliminate all others in a negative test. const allLiteralTypes = enumerateLiteralsForType(evaluator, subtype); if (allLiteralTypes && allLiteralTypes.length > 0) { - return combineTypes(allLiteralTypes.filter((type) => !ClassType.isLiteralValueSame(type, literalType))); + return combineTypes( + allLiteralTypes.filter((type) => + isPrimitiveBackedEnumComparison + ? !isLiteralValueEqualAtRuntime(type, literalType) + : !ClassType.isLiteralValueSame(type, literalType) + ) + ); } return subtype; diff --git a/packages/pyright-internal/src/tests/samples/enumNarrowing1.py b/packages/pyright-internal/src/tests/samples/enumNarrowing1.py new file mode 100644 index 000000000000..e1cd8b75894f --- /dev/null +++ b/packages/pyright-internal/src/tests/samples/enumNarrowing1.py @@ -0,0 +1,79 @@ +# This sample tests type narrowing for equality (== and !=) comparisons +# between IntEnum/StrEnum members and primitive literals or literal unions. + +from enum import IntEnum, IntFlag, StrEnum +from typing import Literal, Union, assert_type + +class Priority(IntEnum): + LOW = 10 + HIGH = 20 + +class Status(StrEnum): + PENDING = "pending" + DONE = "done" + +def test_enum_var_equals_int_literal(p: Priority): + if p == 20: + assert_type(p, Literal[Priority.HIGH]) + else: + assert_type(p, Literal[Priority.LOW]) + +def test_int_var_equals_enum_literal(x: Union[Literal[10], Literal[20], Literal[30]]): + if x == Priority.HIGH: + assert_type(x, Literal[20]) + else: + assert_type(x, Union[Literal[10], Literal[30]]) + +def test_str_var_equals_strenum_literal(s: Union[Literal["pending"], Literal["done"], Literal["failed"]]): + if s == Status.DONE: + assert_type(s, Literal["done"]) + else: + assert_type(s, Union[Literal["pending"], Literal["failed"]]) + +def test_strenum_var_equals_str_literal(st: Status): + if st == "done": + assert_type(st, Literal[Status.DONE]) + else: + assert_type(st, Literal[Status.PENDING]) + + +class Flags(IntFlag): + READ = 1 + WRITE = 2 + + +class CustomInt(int): ... + + +def test_int_flag(f: Flags): + if f == 3: + assert_type(f, Flags) + else: + assert_type(f, Flags) + + +def test_bool_vs_int(x: int, b: bool): + # bool is derived from int, but its literal values are distinct from + # int literal values, so no narrowing should occur here. + if x == True: + assert_type(x, int) + + if b == 1: + assert_type(b, bool) + + +def test_custom_int_subclass(c: CustomInt): + # A custom int subclass is not an enum, so it should not be narrowed + # to a primitive literal type. + if c == 20: + assert_type(c, CustomInt) + + if c == Priority.HIGH: + assert_type(c, CustomInt) + + +def test_plain_enum(e: Priority, y: int): + # An int variable compared against an enum member should retain its + # declared type rather than narrowing to the enum literal. + if y == Priority.HIGH: + assert_type(y, int) diff --git a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts index 3a5a53469bf7..c393ad6e207e 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator3.test.ts +++ b/packages/pyright-internal/src/tests/typeEvaluator3.test.ts @@ -1135,3 +1135,9 @@ test('EnumGenNextValue1', () => { TestUtils.validateResults(analysisResults, 0); }); + +test('EnumNarrowing1', () => { + const analysisResults = TestUtils.typeAnalyzeSampleFiles(['enumNarrowing1.py']); + + TestUtils.validateResults(analysisResults, 0); +});