diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts index 7ee3a18d5b37..364a65b8e1fe 100644 --- a/packages/pyright-internal/src/analyzer/typeGuards.ts +++ b/packages/pyright-internal/src/analyzer/typeGuards.ts @@ -373,6 +373,22 @@ export function getTypeNarrowingCallback( }; }; } + + // Look for X == or X != . + if (isInstantiableClass(rightType)) { + return (type: Type) => { + return { + type: narrowTypeForClassComparison( + evaluator, + type, + rightType, + adjIsPositiveTest, + /* isIsOperator */ false + ), + isIncomplete: !!rightTypeResult.isIncomplete, + }; + }; + } } // Look for X[] == or X[] != @@ -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; } diff --git a/packages/pyright-internal/src/tests/samples/typeGuard4.py b/packages/pyright-internal/src/tests/samples/typeGuard4.py new file mode 100644 index 000000000000..1686d08fd576 --- /dev/null +++ b/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) diff --git a/packages/pyright-internal/src/tests/typeEvaluator6.test.ts b/packages/pyright-internal/src/tests/typeEvaluator6.test.ts index eb1d6a09ab96..ca950fa56e44 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator6.test.ts +++ b/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);