Fix type narrowing for equality comparisons with IntEnum and StrEnum values

[Henry Su (hsusul)]

Summary

Fixes type narrowing for equality (== and !=) comparisons between IntEnum / StrEnum values and primitive literals (or literal unions).

Python Reproduction & Expected Behavior

In Python, IntEnum (derived from int and enum.Enum) and StrEnum (derived from str and enum.Enum) are runtime subtypes of int and str. Equality comparisons such as p == 20 (where p: Priority) or x == Priority.HIGH (where x: int or x: Union[Literal[10], Literal[20]]) evaluate to True at runtime when the values match.

from enum import IntEnum, 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])  # Previously: Priority (un-narrowed)
    else:
        assert_type(p, Literal[Priority.LOW])   # Previously: Priority (un-narrowed)

def test_int_var_equals_enum_literal(x: Union[Literal[10], Literal[20], Literal[30]]):
    if x == Priority.HIGH:
        assert_type(x, Literal[20])              # Previously: Literal[10, 20, 30] (un-narrowed)
    else:
        assert_type(x, Union[Literal[10], Literal[30]])

def test_strenum_var_equals_str_literal(st: Status):
    if st == "done":
        assert_type(st, Literal[Status.DONE])   # Previously: Status (un-narrowed)
    else:
        assert_type(st, Literal[Status.PENDING])

Root Cause & Implementation

1. In narrowTypeForLiteralComparison (packages/pyright-internal/src/analyzer/typeGuards.ts), line 2714 previously checked ClassType.isSameGenericClass(literalType, subtype). When comparing a derived enum class (e.g. Priority or Status) to an underlying primitive type (int or str), isSameGenericClass returned false, causing Pyright to skip literal comparison and return subtype un-narrowed.
   • Fixed by checking isSameOrDerivedClass using the un-literal base class types (ClassType.cloneWithLiteral(..., undefined)) when evaluating equality (!isIsOperator) tests.
2. In ClassType.isLiteralValueSame (packages/pyright-internal/src/analyzer/types.ts), comparing EnumLiteral to number or string returned false.
   • Updated isLiteralValueSame to un-wrap val1.itemType.priv.literalValue (the underlying primitive literal value for IntEnum/StrEnum/ReprEnum members) when comparing with primitive literals.

Regression Coverage

• Added new sample test suite packages/pyright-internal/src/tests/samples/enumNarrowing1.py and registered EnumNarrowing1 test case in typeEvaluator3.test.ts.

Validation Results

• npx jest src/tests/typeEvaluator3.test.ts -t "EnumNarrowing1": PASS
• npx jest src/tests/typeEvaluator1.test.ts ... typeEvaluator8.test.ts: PASS (1208/1208 tests passed)
• Full test suite (npm run test): PASS (64/64 test suites, 2578/2578 tests passed)
• Linting (npm run check): PASS
• Type checking (npm run typecheck): PASS
• Build (npm run build:cli:dev): PASS
• Diff check (git diff --check): PASS (clean)

[Stella Huang (StellaHuang95)]

🔒 Automated review in progress — Stella Huang (@StellaHuang95) is auto-reviewing this PR.

[Stella Huang (StellaHuang95)]

The expanded derived-class gate can narrow ordinary bool/int comparisons to the wrong primitive type, and unmatched enum comparisons can produce invalid primitive literals. These are soundness regressions in shared type-narrowing code.

[Stella Huang (StellaHuang95)]

The derived-class widening introduces unsound narrowing for ordinary bool/int comparisons and for enum comparisons with unmatched literals. These regressions can change variables to incompatible primitive literal types.

[Stella Huang (StellaHuang95)]

The new derived-class path can narrow enum or primitive values to an incompatible literal when no enum member matches, producing unsound types.

[Rich Chiodo (rchiodo)]

🔒 Automated review in progress — Rich Chiodo (@rchiodo) is auto-reviewing this PR.

[Rich Chiodo (rchiodo)]

Issue · Please address or respond

This bidirectional derived-class check admits all primitive subclasses. For example, x: int compared with Priority.HIGH now narrows to Literal[Priority.HIGH] rather than the primitive literal, and a custom int subclass loses its subtype. IntFlag is also affected: it cannot be enumerated here, so a comparison can fall through to a bare primitive literal. Restrict this cross-class path to supported primitive-backed enum comparisons and preserve the reference subtype when no enum member can be selected.

[Rich Chiodo (rchiodo)]

Issue · Please address or respond

This shared literal-identity helper now unwraps every EnumLiteral, including ordinary Enum members whose runtime equality does not match their underlying value. Please restrict this behavior to enum classes with primitive runtime equality semantics, or keep it in an equality-specific helper.

[Rich Chiodo (rchiodo)]

Issue · Please address or respond

The symmetric derived-class check is not enum-specific. For example, x: int; if x == True can now narrow x to Literal[True], even though the matching runtime value may be the integer 1. Restrict cross-class compatibility to the supported enum/primitive pairs and add a bool-versus-int regression test.

[Rich Chiodo (rchiodo)]

Approved via Review Center.

[Heejae Chang (heejaechang)]

Approved via Review Center.