Fixed bug that results in a false negative when attempting to assign an IntEnum or StrEnum literal to the Literal type corresponding to its value type. This addresses #10552. (#10558)

Author: erictraut

packages/pyright-internal/src/analyzer/checker.ts

@@ -129,6 +129,7 @@ import {
     TypeResult,
 } from './typeEvaluatorTypes';
 import {
+    enumerateLiteralsForType,
     getElementTypeForContainerNarrowing,
     getIsInstanceClassTypes,
     narrowTypeForContainerElementType,
@@ -2142,8 +2143,8 @@ export class Checker extends ParseTreeWalker {
             rightExpression = rightExpression.d.leftExpr;
         }
 
-        const leftType = this._evaluator.getType(node.d.leftExpr);
-        const rightType = this._evaluator.getType(rightExpression);
+        let leftType = this._evaluator.getType(node.d.leftExpr);
+        let rightType = this._evaluator.getType(rightExpression);
 
         if (!leftType || !rightType) {
             return;
@@ -2159,6 +2160,44 @@ export class Checker extends ParseTreeWalker {
                 : LocMessage.comparisonAlwaysTrue();
         };
 
+        const replaceEnumTypeWithLiteralValue = (type: Type) => {
+            return mapSubtypes(type, (subtype) => {
+                if (
+                    !isClassInstance(subtype) ||
+                    !ClassType.isEnumClass(subtype) ||
+                    !subtype.shared.mro.some(
+                        (base) => isClass(base) && ClassType.isBuiltIn(base, ['int', 'str', 'bytes'])
+                    )
+                ) {
+                    return subtype;
+                }
+
+                // If this is an enum literal, replace it with its literal value.
+                if (subtype.priv.literalValue instanceof EnumLiteral) {
+                    return subtype.priv.literalValue.itemType;
+                }
+
+                // If this is an enum class, replace it with the type of its members.
+                const literalValues = enumerateLiteralsForType(this._evaluator, subtype);
+                if (literalValues && literalValues.length > 0) {
+                    return combineTypes(
+                        literalValues.map((literalClass) => {
+                            const literalValue = literalClass.priv.literalValue;
+                            assert(literalValue instanceof EnumLiteral);
+                            return literalValue.itemType;
+                        })
+                    );
+                }
+
+                return subtype;
+            });
+        };
+
+        // Handle enum literals that are assignable to another (non-Enum) literal.
+        // This can happen for IntEnum and StrEnum members.
+        leftType = replaceEnumTypeWithLiteralValue(leftType);
+        rightType = replaceEnumTypeWithLiteralValue(rightType);
+
         // Check for the special case where the LHS and RHS are both literals.
         if (isLiteralTypeOrUnion(rightType) && isLiteralTypeOrUnion(leftType)) {
             if (
@@ -2171,13 +2210,13 @@ export class Checker extends ParseTreeWalker {
                 let isPossiblyTrue = false;
 
                 doForEachSubtype(leftType, (leftSubtype) => {
-                    if (this._evaluator.assignType(rightType, leftSubtype)) {
+                    if (this._evaluator.assignType(rightType!, leftSubtype)) {
                         isPossiblyTrue = true;
                     }
                 });
 
                 doForEachSubtype(rightType, (rightSubtype) => {
-                    if (this._evaluator.assignType(leftType, rightSubtype)) {
+                    if (this._evaluator.assignType(leftType!, rightSubtype)) {
                         isPossiblyTrue = true;
                     }
                 });
@@ -2201,7 +2240,7 @@ export class Checker extends ParseTreeWalker {
                     return;
                 }
 
-                this._evaluator.mapSubtypesExpandTypeVars(rightType, {}, (rightSubtype) => {
+                this._evaluator.mapSubtypesExpandTypeVars(rightType!, {}, (rightSubtype) => {
                     if (isComparable) {
                         return;
                     }

packages/pyright-internal/src/analyzer/typeEvaluator.ts

@@ -24903,21 +24903,6 @@ export function createTypeEvaluator(
                     );
                 }
 
-                // Handle enum literals that are assignable to another (non-Enum) literal.
-                // This can happen for IntEnum and StrEnum members.
-                if (
-                    ClassType.isEnumClass(concreteSrcType) &&
-                    concreteSrcType.priv.literalValue instanceof EnumLiteral &&
-                    concreteSrcType.shared.mro.some(
-                        (base) => isClass(base) && ClassType.isBuiltIn(base, ['int', 'str', 'bytes'])
-                    ) &&
-                    isClassInstance(concreteSrcType.priv.literalValue.itemType) &&
-                    isLiteralType(concreteSrcType.priv.literalValue.itemType) &&
-                    assignType(destType, concreteSrcType.priv.literalValue.itemType)
-                ) {
-                    return true;
-                }
-
                 if (
                     destType.priv.literalValue !== undefined &&
                     ClassType.isSameGenericClass(destType, concreteSrcType)

packages/pyright-internal/src/tests/samples/comparison2.py

@@ -2,6 +2,7 @@
 # when applied to functions that appear within a conditional expression.
 
 
+from enum import Enum
 from typing import Any, Callable, Coroutine, Protocol
 from dataclasses import dataclass
 
@@ -140,3 +141,18 @@ def func11(a: A, b: SupportsBool, c: object):
 
 def func12(a: object, b: Callable[..., int]) -> bool:
     return a is b
+
+
+class IntVal(int, Enum):
+    one = 1
+    two = 2
+    three = 3
+
+
+def func13(x: IntVal):
+    if x == 1:
+        pass
+
+    # This should generate an error if reportUnnecessaryComparison is enabled.
+    if x == 4:
+        pass

packages/pyright-internal/src/tests/samples/enum13.py

@@ -9,31 +9,31 @@ class IntEnum1(IntEnum):
     MEMBER_2 = 2
 
 
-i1: Literal[1] = IntEnum1.MEMBER_1
+i1: Literal[1] = IntEnum1.MEMBER_1.value
 
 # This should generate an error.
-i2: Literal[1] = IntEnum1.MEMBER_2
+i2: Literal[1] = IntEnum1.MEMBER_2.value
 
 
 class StrEnum1(StrEnum):
     MEMBER_1 = "a"
     MEMBER_2 = "b"
 
 
-s1: Literal["a"] = StrEnum1.MEMBER_1
+s1: Literal["a"] = StrEnum1.MEMBER_1.value
 
 # This should generate an error.
-s2: Literal["b"] = StrEnum1.MEMBER_1
+s2: Literal["b"] = StrEnum1.MEMBER_1.value
 
-s3: LiteralString = StrEnum1.MEMBER_1
+s3: LiteralString = StrEnum1.MEMBER_1.value
 
 
 class BytesEnum(bytes, ReprEnum):
     MEMBER_1 = b"1"
     MEMBER_2 = b"2"
 
 
-b1: Literal[b"1"] = BytesEnum.MEMBER_1
+b1: Literal[b"1"] = BytesEnum.MEMBER_1.value
 
 # This should generate an error.
-b2: Literal[b"2"] = BytesEnum.MEMBER_1
+b2: Literal[b"2"] = BytesEnum.MEMBER_1.value

packages/pyright-internal/src/tests/typeEvaluator6.test.ts

@@ -629,7 +629,7 @@ test('Comparison2', () => {
 
     configOptions.diagnosticRuleSet.reportUnnecessaryComparison = 'error';
     const analysisResults2 = TestUtils.typeAnalyzeSampleFiles(['comparison2.py'], configOptions);
-    TestUtils.validateResults(analysisResults2, 17);
+    TestUtils.validateResults(analysisResults2, 18);
 });
 
 test('EmptyContainers1', () => {