From bfc7466bb9ad037cbb11ec84a572006ae75d9e30 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Sun, 9 Aug 2026 23:30:44 -0500 Subject: [PATCH] Fix TypeGuard and TypeIs narrowing for functions with union return types --- .../pyright-internal/src/analyzer/checker.ts | 69 +++++++++++-------- .../src/analyzer/typeGuards.ts | 59 +++++++++++++--- .../src/tests/samples/typeIs5.py | 61 ++++++++++++++++ .../src/tests/typeEvaluator6.test.ts | 5 ++ 4 files changed, 156 insertions(+), 38 deletions(-) create mode 100644 packages/pyright-internal/src/tests/samples/typeIs5.py diff --git a/packages/pyright-internal/src/analyzer/checker.ts b/packages/pyright-internal/src/analyzer/checker.ts index 58db6863cc96..8f7f90fe1c76 100644 --- a/packages/pyright-internal/src/analyzer/checker.ts +++ b/packages/pyright-internal/src/analyzer/checker.ts @@ -4669,14 +4669,19 @@ export class Checker extends ParseTreeWalker { return; } - if (!isClassInstance(returnType) || !returnType.priv.typeArgs || returnType.priv.typeArgs.length < 1) { - return; - } - - const isTypeGuard = ClassType.isBuiltIn(returnType, 'TypeGuard'); - const isTypeIs = ClassType.isBuiltIn(returnType, 'TypeIs'); + const guardSubtypes: ClassType[] = []; + doForEachSubtype(returnType, (subtype) => { + if ( + isClassInstance(subtype) && + (ClassType.isBuiltIn(subtype, 'TypeGuard') || ClassType.isBuiltIn(subtype, 'TypeIs')) && + subtype.priv.typeArgs && + subtype.priv.typeArgs.length >= 1 + ) { + guardSubtypes.push(subtype); + } + }); - if (!isTypeGuard && !isTypeIs) { + if (guardSubtypes.length === 0) { return; } @@ -4700,32 +4705,36 @@ export class Checker extends ParseTreeWalker { ); } - if (isTypeIs) { - const scopeIds = getTypeVarScopeIds(functionType); - const narrowedType = returnType.priv.typeArgs[0]; - let typeGuardType = makeTypeVarsBound(narrowedType, scopeIds); - typeGuardType = TypeBase.cloneWithTypeForm(typeGuardType, typeGuardType); + const scopeIds = getTypeVarScopeIds(functionType); - // Determine the type of the first parameter. - const paramIndex = isMethod && !FunctionType.isStaticMethod(functionType) ? 1 : 0; - if (paramIndex >= functionType.shared.parameters.length) { - return; - } + // Determine the type of the first parameter. + const paramIndex = isMethod && !FunctionType.isStaticMethod(functionType) ? 1 : 0; + if (paramIndex >= functionType.shared.parameters.length) { + return; + } - const paramType = makeTypeVarsBound(FunctionType.getParamType(functionType, paramIndex), scopeIds); + const paramType = makeTypeVarsBound(FunctionType.getParamType(functionType, paramIndex), scopeIds); - // Verify that the typeGuardType is a narrower type than the paramType. - if (!this._evaluator.assignType(paramType, typeGuardType)) { - const returnAnnotation = node.d.returnAnnotation || node.d.funcAnnotationComment?.d.returnAnnotation; - if (returnAnnotation) { - this._evaluator.addDiagnostic( - DiagnosticRule.reportGeneralTypeIssues, - LocMessage.typeIsReturnType().format({ - type: this._evaluator.printType(paramType), - returnType: this._evaluator.printType(narrowedType), - }), - returnAnnotation - ); + for (const guardSubtype of guardSubtypes) { + if (ClassType.isBuiltIn(guardSubtype, 'TypeIs')) { + const narrowedType = guardSubtype.priv.typeArgs![0]; + let typeGuardType = makeTypeVarsBound(narrowedType, scopeIds); + typeGuardType = TypeBase.cloneWithTypeForm(typeGuardType, typeGuardType); + + // Verify that the typeGuardType is a narrower type than the paramType. + if (!this._evaluator.assignType(paramType, typeGuardType)) { + const returnAnnotation = + node.d.returnAnnotation || node.d.funcAnnotationComment?.d.returnAnnotation; + if (returnAnnotation) { + this._evaluator.addDiagnostic( + DiagnosticRule.reportGeneralTypeIssues, + LocMessage.typeIsReturnType().format({ + type: this._evaluator.printType(paramType), + returnType: this._evaluator.printType(narrowedType), + }), + returnAnnotation + ); + } } } } diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts index 7ee3a18d5b37..35d44b5ad90a 100644 --- a/packages/pyright-internal/src/analyzer/typeGuards.ts +++ b/packages/pyright-internal/src/analyzer/typeGuards.ts @@ -55,6 +55,7 @@ import { isTypeSame, isTypeVar, isUnpackedTypeVarTuple, + isUnion, maxTypeRecursionCount, OverloadedType, TupleTypeArg, @@ -709,11 +710,23 @@ export function getTypeNarrowingCallback( let isPossiblyTypeGuard = false; const isFunctionReturnTypeGuard = (type: FunctionType) => { - return ( - type.shared.declaredReturnType && - isClassInstance(type.shared.declaredReturnType) && - ClassType.isBuiltIn(type.shared.declaredReturnType, ['TypeGuard', 'TypeIs']) - ); + const returnType = type.shared.declaredReturnType; + if (!returnType) { + return false; + } + if (isClassInstance(returnType)) { + return ClassType.isBuiltIn(returnType, ['TypeGuard', 'TypeIs']); + } + if (isUnion(returnType)) { + let isAllGuards = true; + doForEachSubtype(returnType, (subtype) => { + if (!isClassInstance(subtype) || !ClassType.isBuiltIn(subtype, ['TypeGuard', 'TypeIs'])) { + isAllGuards = false; + } + }); + return isAllGuards; + } + return false; }; const callTypeResult = evaluator.getTypeOfExpression( @@ -738,14 +751,44 @@ export function getTypeNarrowingCallback( const functionReturnTypeResult = evaluator.getTypeOfExpression(testExpression); const functionReturnType = functionReturnTypeResult.type; + let typeGuardType: Type | undefined; + let isStrictTypeGuard = false; + if ( isClassInstance(functionReturnType) && ClassType.isBuiltIn(functionReturnType, ['TypeGuard', 'TypeIs']) && functionReturnType.priv.typeArgs && functionReturnType.priv.typeArgs.length > 0 ) { - const isStrictTypeGuard = ClassType.isBuiltIn(functionReturnType, 'TypeIs'); - const typeGuardType = functionReturnType.priv.typeArgs[0]; + isStrictTypeGuard = ClassType.isBuiltIn(functionReturnType, 'TypeIs'); + typeGuardType = functionReturnType.priv.typeArgs[0]; + } else if (isUnion(functionReturnType)) { + const typeGuardSubtypes: ClassType[] = []; + let isAllGuards = true; + + doForEachSubtype(functionReturnType, (subtype) => { + if ( + isClassInstance(subtype) && + ClassType.isBuiltIn(subtype, ['TypeGuard', 'TypeIs']) && + subtype.priv.typeArgs && + subtype.priv.typeArgs.length > 0 + ) { + typeGuardSubtypes.push(subtype); + } else { + isAllGuards = false; + } + }); + + if (isAllGuards && typeGuardSubtypes.length > 0) { + // A union of type guards cannot be strict in the negative case because at runtime + // only one arm of the overload/union is selected. Treating it as strict would + // unsoundly eliminate types in the negative branch. + isStrictTypeGuard = false; + typeGuardType = combineTypes(typeGuardSubtypes.map((subtype) => subtype.priv.typeArgs![0])); + } + } + + if (typeGuardType) { const isIncomplete = !!callTypeResult.isIncomplete || !!functionReturnTypeResult.isIncomplete; return (type: Type) => { @@ -753,7 +796,7 @@ export function getTypeNarrowingCallback( type: narrowTypeForUserDefinedTypeGuard( evaluator, type, - typeGuardType, + typeGuardType!, isPositiveTest, isStrictTypeGuard, testExpression diff --git a/packages/pyright-internal/src/tests/samples/typeIs5.py b/packages/pyright-internal/src/tests/samples/typeIs5.py new file mode 100644 index 000000000000..312c11449acd --- /dev/null +++ b/packages/pyright-internal/src/tests/samples/typeIs5.py @@ -0,0 +1,61 @@ +# This sample tests user-defined TypeIs and TypeGuard functions whose return type +# is a union of TypeIs or TypeGuard instances. + +from typing import TypeGuard, TypeIs, assert_type, overload + + +def check_single(val: object) -> TypeIs[int] | TypeIs[str]: + return isinstance(val, (int, str)) + + +# This should generate an error because "int" is not a subtype of "str". +def invalid_typeis_union(val: str) -> TypeIs[int] | TypeIs[str]: # pyright: ignore[reportGeneralTypeIssues] + return False + + +@overload +def check_overload(val: object, target: type[int]) -> TypeIs[int]: ... +@overload +def check_overload(val: object, target: type[str]) -> TypeIs[str]: ... + + +def check_overload(val: object, target: type) -> bool: + return isinstance(val, target) + + +def check_mixed(val: object) -> TypeIs[int] | TypeGuard[str]: + return isinstance(val, (int, str)) + + +def check_nonguard(val: object) -> TypeIs[int] | None: + return isinstance(val, int) if val else None + + +def test_single(x: object): + if check_single(x): + assert_type(x, int | str) + + +def test_overload_positive(x: object, target: type[int] | type[str]): + if check_overload(x, target): + assert_type(x, int | str) + + +def test_overload_negative(x: int | str | bytes, target: type[int] | type[str]): + if check_overload(x, target): + assert_type(x, int | str) + else: + # A union of type guards is non-strict in the negative case to prevent + # unsound type elimination when only one overload/arm applies at runtime. + assert_type(x, int | str | bytes) + + +def test_mixed(x: object): + if check_mixed(x): + assert_type(x, int | str) + + +def test_nonguard(x: object): + if check_nonguard(x): + # Non-guard members in the return type union cause the type guard to be rejected. + 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..81545ff2429b 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator6.test.ts +++ b/packages/pyright-internal/src/tests/typeEvaluator6.test.ts @@ -159,6 +159,11 @@ test('TypeIs4', () => { TestUtils.validateResults(analysisResults, 0); }); +test('TypeIs5', () => { + const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeIs5.py']); + TestUtils.validateResults(analysisResults, 0); +}); + test('Never1', () => { const analysisResults = TestUtils.typeAnalyzeSampleFiles(['never1.py']);