Fixed a bug that results in incorrect type narrowing for truthy/falsy type guards consisting of an Enum member when the enum class derives from ReprEnum. In this case, the "magic" of the enum implementation forwards the `__bool__` call to the underlying value. This addresses #8658. (#8662)

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

@@ -72,6 +72,7 @@ export function createEnumType(
     argList: Arg[]
 ): ClassType | undefined {
     const fileInfo = getFileInfo(errorNode);
+    const isReprEnum = isReprEnumClass(enumClass);
 
     if (argList.length === 0) {
         return undefined;
@@ -152,7 +153,13 @@ export function createEnumType(
 
             const valueType = ClassType.cloneWithLiteral(ClassType.cloneAsInstance(intClassType), index + 1);
 
-            const enumLiteral = new EnumLiteral(classType.shared.fullName, classType.shared.name, entryName, valueType);
+            const enumLiteral = new EnumLiteral(
+                classType.shared.fullName,
+                classType.shared.name,
+                entryName,
+                valueType,
+                isReprEnum
+            );
 
             const newSymbol = Symbol.createWithType(
                 SymbolFlags.ClassMember,
@@ -220,7 +227,13 @@ export function createEnumType(
 
             const entryName = nameNode.d.strings[0].d.value;
 
-            const enumLiteral = new EnumLiteral(classType.shared.fullName, classType.shared.name, entryName, valueType);
+            const enumLiteral = new EnumLiteral(
+                classType.shared.fullName,
+                classType.shared.name,
+                entryName,
+                valueType,
+                isReprEnum
+            );
 
             const newSymbol = Symbol.createWithType(
                 SymbolFlags.ClassMember,
@@ -255,7 +268,13 @@ export function createEnumType(
             }
 
             const entryName = nameNode.d.strings[0].d.value;
-            const enumLiteral = new EnumLiteral(classType.shared.fullName, classType.shared.name, entryName, valueType);
+            const enumLiteral = new EnumLiteral(
+                classType.shared.fullName,
+                classType.shared.name,
+                entryName,
+                valueType,
+                isReprEnum
+            );
 
             const newSymbol = Symbol.createWithType(
                 SymbolFlags.ClassMember,
@@ -475,7 +494,8 @@ export function transformTypeForEnumMember(
         memberInfo.classType.shared.fullName,
         memberInfo.classType.shared.name,
         memberName,
-        valueType
+        valueType,
+        isReprEnumClass(classType)
     );
 
     return ClassType.cloneAsInstance(ClassType.cloneWithLiteral(memberInfo.classType, enumLiteral));
@@ -691,3 +711,7 @@ export function getEnumAutoValueType(evaluator: TypeEvaluator, node: ExpressionN
 
     return evaluator.getBuiltInObject(node, 'int');
 }
+
+function isReprEnumClass(enumClass: ClassType) {
+    return enumClass.shared.mro.some((mroClass) => isClass(mroClass) && ClassType.isBuiltIn(mroClass, 'ReprEnum'));
+}

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

@@ -1830,11 +1830,18 @@ export function createTypeEvaluator(
                 }
 
                 // Check for bool, int, str and bytes literals that are never falsy.
-                if (
-                    type.priv.literalValue !== undefined &&
-                    ClassType.isBuiltIn(type, ['bool', 'int', 'str', 'bytes'])
-                ) {
-                    return !type.priv.literalValue || type.priv.literalValue === BigInt(0);
+                if (type.priv.literalValue !== undefined) {
+                    if (ClassType.isBuiltIn(type, ['bool', 'int', 'str', 'bytes'])) {
+                        return !type.priv.literalValue || type.priv.literalValue === BigInt(0);
+                    }
+
+                    if (type.priv.literalValue instanceof EnumLiteral) {
+                        // Does the Enum class forward the truthiness check to the
+                        // underlying member type?
+                        if (type.priv.literalValue.isReprEnum) {
+                            return canBeFalsy(type.priv.literalValue.itemType, recursionCount);
+                        }
+                    }
                 }
 
                 // If this is a protocol class, don't make any assumptions about the absence
@@ -1921,11 +1928,18 @@ export function createTypeEvaluator(
                 }
 
                 // Check for bool, int, str and bytes literals that are never falsy.
-                if (
-                    type.priv.literalValue !== undefined &&
-                    ClassType.isBuiltIn(type, ['bool', 'int', 'str', 'bytes'])
-                ) {
-                    return !!type.priv.literalValue && type.priv.literalValue !== BigInt(0);
+                if (type.priv.literalValue !== undefined) {
+                    if (ClassType.isBuiltIn(type, ['bool', 'int', 'str', 'bytes'])) {
+                        return !!type.priv.literalValue && type.priv.literalValue !== BigInt(0);
+                    }
+
+                    if (type.priv.literalValue instanceof EnumLiteral) {
+                        // Does the Enum class forward the truthiness check to the
+                        // underlying member type?
+                        if (type.priv.literalValue.isReprEnum) {
+                            return canBeTruthy(type.priv.literalValue.itemType, recursionCount);
+                        }
+                    }
                 }
 
                 // If this is a protocol class, don't make any assumptions about the absence

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

@@ -90,7 +90,8 @@ export class EnumLiteral {
         public classFullName: string,
         public className: string,
         public itemName: string,
-        public itemType: Type
+        public itemType: Type,
+        public isReprEnum: boolean
     ) {}
 
     getName() {

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

@@ -184,3 +184,15 @@ def func17(x: Enum3):
         reveal_type(x, expected_text="Enum3")
     else:
         reveal_type(x, expected_text="Enum3")
+
+
+def func18(x: Literal[Enum3.A], y: Literal[Enum3.B]):
+    if x:
+        reveal_type(x, expected_text="Never")
+    else:
+        reveal_type(x, expected_text="Literal[Enum3.A]")
+
+    if y:
+        reveal_type(y, expected_text="Literal[Enum3.B]")
+    else:
+        reveal_type(y, expected_text="Never")