Fix type narrowing for class equality comparisons

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

@@ -373,6 +373,22 @@ export function getTypeNarrowingCallback(
                             };
                         };
                     }
+
+                    // Look for X == <class> or X != <class>.
+                    if (isInstantiableClass(rightType)) {
+                        return (type: Type) => {
+                            return {
+                                type: narrowTypeForClassComparison(
+                                    evaluator,
+                                    type,
+                                    rightType,
+                                    adjIsPositiveTest,
+                                    /* isIsOperator */ false
+                                ),
+                                isIncomplete: !!rightTypeResult.isIncomplete,
+                            };
+                        };
+                    }
                 }
 
                 // Look for X[<literal>] == <literal> or X[<literal>] != <literal>
@@ -2585,17 +2601,47 @@ function narrowTypeForTypeIs(evaluator: TypeEvaluator, type: Type, classTypes: C
     return combineTypes(typesToCombine);
 }
 
+function hasCustomEqualityMetaclass(classType: ClassType): boolean {
+    const metaclass = classType.shared.effectiveMetaclass;
+    if (metaclass && isClass(metaclass)) {
+        if (
+            lookUpClassMember(
+                metaclass,
+                '__eq__',
+                MemberAccessFlags.SkipTypeBaseClass | MemberAccessFlags.SkipObjectBaseClass
+            ) ||
+            lookUpClassMember(
+                metaclass,
+                '__ne__',
+                MemberAccessFlags.SkipTypeBaseClass | MemberAccessFlags.SkipObjectBaseClass
+            )
+        ) {
+            return true;
+        }
+    }
+    return false;
+}
+
 // Attempts to narrow a type based on a comparison with a class using "is" or
-// "is not". This pattern is sometimes used for sentinels.
+// "is not", or "==" or "!=".
 function narrowTypeForClassComparison(
     evaluator: TypeEvaluator,
     referenceType: Type,
     classType: ClassType,
-    isPositiveTest: boolean
+    isPositiveTest: boolean,
+    isIsOperator = true
 ): Type {
+    if (!isIsOperator && hasCustomEqualityMetaclass(classType)) {
+        return referenceType;
+    }
+
     return mapSubtypes(referenceType, (subtype) => {
         let concreteSubtype = evaluator.makeTopLevelTypeVarsConcrete(subtype);
 
+        if (!isIsOperator && isInstantiableClass(concreteSubtype) && hasCustomEqualityMetaclass(concreteSubtype)) {
+            return subtype;
+        }
+
         if (isPositiveTest) {
             if (
                 isClassInstance(concreteSubtype) &&
@@ -2614,6 +2660,9 @@ function narrowTypeForClassComparison(
 
             if (isClass(concreteSubtype)) {
                 if (TypeBase.isInstance(concreteSubtype)) {
+                    if (!isIsOperator) {
+                        return subtype;
+                    }
                     return ClassType.isBuiltIn(concreteSubtype, 'object') ? classType : undefined;
                 }
 

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

@@ -0,0 +1,51 @@
+# This sample tests type narrowing when comparing class types
+# with equality (== and !=) operators against class objects.
+
+from typing import TypeVar, assert_type, final
+
+class Base: pass
+class Sub1(Base): pass
+
+@final
+class Sub2(Base): pass
+
+T = TypeVar("T", bound=Base)
+
+def test_eq_concrete(cls: type[Base]) -> type[Sub1]:
+    if cls == Sub1:
+        assert_type(cls, type[Sub1])
+        return cls
+    raise ValueError()
+
+def test_neq_concrete(cls: type[Sub1] | type[Sub2]):
+    if cls != Sub2:
+        assert_type(cls, type[Sub1])
+
+def test_neq_non_final(cls: type[Sub1] | type[Sub2]):
+    if cls != Sub1:
+        assert_type(cls, type[Sub1] | type[Sub2])
+
+def test_eq_typevar(cls: type[T]) -> type[Sub1]:
+    if cls == Sub1:
+        assert_type(cls, type[Sub1])
+        return cls
+    raise ValueError()
+
+class CustomMeta(type):
+    def __eq__(cls, other: object) -> bool:
+        return True
+
+class Custom1(metaclass=CustomMeta): pass
+class Custom2(metaclass=CustomMeta): pass
+
+def test_eq_custom_meta(cls: type[Custom1] | type[Custom2]):
+    if cls == Custom1:
+        assert_type(cls, type[Custom1] | type[Custom2])
+
+class EqualityDummy:
+    def __eq__(self, other: object) -> bool:
+        return True
+
+def test_eq_instance_object(x: object):
+    if x == Sub1:
+        assert_type(x, object)

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

@@ -139,6 +139,11 @@ test('TypeGuard3', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('TypeGuard4', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeGuard4.py']);
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('TypeIs1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeIs1.py']);
     TestUtils.validateResults(analysisResults, 2);