diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts index 9c5e69bc3612..5d9b0b263293 100644 --- a/packages/pyright-internal/src/analyzer/typeGuards.ts +++ b/packages/pyright-internal/src/analyzer/typeGuards.ts @@ -13,6 +13,7 @@ import { assert } from '../common/debug'; import { ArgCategory, AssignmentExpressionNode, + CallNode, ExpressionNode, isExpressionNode, NameNode, @@ -209,45 +210,83 @@ export function getTypeNarrowingCallback( } } - // Look for "type(X) is Y", "type(X) is not Y", "type(X) == Y" or "type(X) != Y". + // Look for "type(X) is Y", "type(X) is not Y", "type(X) == Y" or "type(X) != Y", + // as well as reverse forms "Y is type(X)", "Y is not type(X)", "Y == type(X)", "Y != type(X)". + let typeCallNode: CallNode | undefined; + let otherExpr: ExpressionNode | undefined; + if (testExpression.d.leftExpr.nodeType === ParseNodeType.Call) { - if ( - testExpression.d.leftExpr.d.args.length === 1 && - testExpression.d.leftExpr.d.args[0].d.argCategory === ArgCategory.Simple - ) { - const arg0Expr = testExpression.d.leftExpr.d.args[0].d.valueExpr; - if (isMatchingExpressionOrWalrusRhs(evaluator, reference, arg0Expr)) { - const callType = evaluator.getTypeOfExpression( - testExpression.d.leftExpr.d.leftExpr, - EvalFlags.CallBaseDefaults - ).type; - - if (isInstantiableClass(callType) && ClassType.isBuiltIn(callType, 'type')) { - const rhsResult = evaluator.getTypeOfExpression(testExpression.d.rightExpr); - const classTypes: ClassType[] = []; - let isClassType = true; - - evaluator.mapSubtypesExpandTypeVars( - rhsResult.type, - /* options */ undefined, - (expandedSubtype) => { - if (isInstantiableClass(expandedSubtype)) { - classTypes.push(expandedSubtype); - } else { - isClassType = false; + typeCallNode = testExpression.d.leftExpr; + otherExpr = testExpression.d.rightExpr; + } else if (testExpression.d.rightExpr.nodeType === ParseNodeType.Call) { + typeCallNode = testExpression.d.rightExpr; + otherExpr = testExpression.d.leftExpr; + } + + if ( + typeCallNode && + otherExpr && + typeCallNode.d.args.length === 1 && + typeCallNode.d.args[0].d.argCategory === ArgCategory.Simple + ) { + const arg0Expr = typeCallNode.d.args[0].d.valueExpr; + if (isMatchingExpressionOrWalrusRhs(evaluator, reference, arg0Expr)) { + const callType = evaluator.getTypeOfExpression( + typeCallNode.d.leftExpr, + EvalFlags.CallBaseDefaults + ).type; + + if (isInstantiableClass(callType) && ClassType.isBuiltIn(callType, 'type')) { + const otherResult = evaluator.getTypeOfExpression(otherExpr); + const classTypes: ClassType[] = []; + let isClassType = true; + + evaluator.mapSubtypesExpandTypeVars( + otherResult.type, + /* options */ undefined, + (expandedSubtype) => { + let instantiable: ClassType | undefined; + let isTypeParam = false; + + if (isClass(expandedSubtype)) { + if ( + ClassType.isBuiltIn(expandedSubtype, 'type') && + expandedSubtype.priv.typeArgs && + expandedSubtype.priv.typeArgs.length > 0 + ) { + isTypeParam = true; + const extracted = convertToInstantiable( + expandedSubtype.priv.typeArgs[0], + /* includeSubclasses */ true + ); + if (isInstantiableClass(extracted)) { + instantiable = extracted; + } + } else if (isInstantiableClass(expandedSubtype)) { + instantiable = expandedSubtype; } - return undefined; } - ); - if (isClassType && classTypes.length > 0) { - return (type: Type) => { - return { - type: narrowTypeForTypeIs(evaluator, type, classTypes, adjIsPositiveTest), - isIncomplete: !!rhsResult.isIncomplete, - }; - }; + if (instantiable && !ClassType.isBuiltIn(instantiable, 'object')) { + classTypes.push( + isTypeParam + ? instantiable + : ClassType.cloneIncludeSubclasses(instantiable, false) + ); + } else { + isClassType = false; + } + return undefined; } + ); + + if (isClassType && classTypes.length > 0) { + return (type: Type) => { + return { + type: narrowTypeForTypeIs(evaluator, type, classTypes, adjIsPositiveTest), + isIncomplete: !!otherResult.isIncomplete, + }; + }; } } } @@ -2538,17 +2577,20 @@ function narrowTypeForTypeIs(evaluator: TypeEvaluator, type: Type, classTypes: C /* options */ undefined, (subtype: Type, unexpandedSubtype: Type) => { if (isClassInstance(subtype)) { - const matches = ClassType.isDerivedFrom(classType, ClassType.cloneAsInstantiable(subtype)); + const instantiableSubtype = ClassType.cloneAsInstantiable(subtype); + const matches = ClassType.isDerivedFrom(classType, instantiableSubtype); + const isSubclass = ClassType.isDerivedFrom(instantiableSubtype, classType); + if (isPositiveTest) { if (matches) { - if (ClassType.isSameGenericClass(ClassType.cloneAsInstantiable(subtype), classType)) { + if (ClassType.isSameGenericClass(instantiableSubtype, classType)) { return addConditionToType(subtype, getTypeCondition(classType)); } return addConditionToType(ClassType.cloneAsInstance(classType), subtype.props?.condition); } - if (!classType.priv.includeSubclasses) { + if (!classType.priv.includeSubclasses || !isSubclass) { return undefined; } @@ -2560,7 +2602,7 @@ function narrowTypeForTypeIs(evaluator: TypeEvaluator, type: Type, classTypes: C if (!classType.priv.includeSubclasses) { // If the class if marked final and it matches, then // we can eliminate it in the negative case. - if (matches && ClassType.isFinal(subtype)) { + if (matches && ClassType.isFinal(instantiableSubtype)) { return undefined; } diff --git a/packages/pyright-internal/src/tests/samples/typeIs6.py b/packages/pyright-internal/src/tests/samples/typeIs6.py new file mode 100644 index 000000000000..2a5c9efeca9b --- /dev/null +++ b/packages/pyright-internal/src/tests/samples/typeIs6.py @@ -0,0 +1,76 @@ +# This sample tests type narrowing for type(x) checks. + +from typing import final + +@final +class FinalA: pass + +@final +class FinalB: pass + +def func1(x: int | str, cls: type[int]): + if type(x) is cls: + reveal_type(x, expected_text="int") + else: + reveal_type(x, expected_text="int | str") + +def func2(x: int | str, y: int): + if type(x) is type(y): + reveal_type(x, expected_text="int") + else: + reveal_type(x, expected_text="int | str") + +def func3(x: int | str, cls: type[int]): + if cls is type(x): + reveal_type(x, expected_text="int") + else: + reveal_type(x, expected_text="int | str") + +def func4(x: int | str): + if int is type(x): + reveal_type(x, expected_text="int") + else: + reveal_type(x, expected_text="int | str") + +def func5(x: int | str, cls: type[int]): + if type(x) == cls: + reveal_type(x, expected_text="int") + else: + reveal_type(x, expected_text="int | str") + +def func6(x: FinalA | FinalB, cls: type[FinalA]): + if type(x) is cls: + reveal_type(x, expected_text="FinalA") + else: + reveal_type(x, expected_text="FinalB") + +def func7(x: int | str, cls: type[int] | type[str]): + if type(x) is cls: + reveal_type(x, expected_text="int | str") + +class Base: pass + +@final +class FinalSub(Base): pass + +def func8(x: Base): + if type(x) is not FinalSub: + reveal_type(x, expected_text="Base") + else: + reveal_type(x, expected_text="FinalSub") + +class Index: pass + +class MultiIndex(Index): pass + +def test_spark_regression(self_val: MultiIndex, other: Index): + if type(self_val) is not type(other): + pass + reveal_type(self_val, expected_text="MultiIndex") + +def test_direct_class_vs_type_param(x: Base, cls: type[Base]): + if type(x) is not Base: + reveal_type(x, expected_text="Base") + + if type(x) is not cls: + reveal_type(x, expected_text="Base") diff --git a/packages/pyright-internal/src/tests/typeEvaluator6.test.ts b/packages/pyright-internal/src/tests/typeEvaluator6.test.ts index 7e94cc45f778..fff78edf153a 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator6.test.ts +++ b/packages/pyright-internal/src/tests/typeEvaluator6.test.ts @@ -173,6 +173,11 @@ test('TypeIs4', () => { TestUtils.validateResults(analysisResults, 0); }); +test('TypeIs6', () => { + const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeIs6.py']); + TestUtils.validateResults(analysisResults, 0); +}); + test('Never1', () => { const analysisResults = TestUtils.typeAnalyzeSampleFiles(['never1.py']);