Fix type narrowing for equality comparisons with IntEnum and StrEnum values

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,14 +2804,38 @@ 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;
             }
 
             // If we're able to enumerate all possible literal values
             // (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;

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)

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);
+});