diff --git a/docs/type-concepts.md b/docs/type-concepts.md index adfe0bba6..6660f1557 100644 --- a/docs/type-concepts.md +++ b/docs/type-concepts.md @@ -168,6 +168,8 @@ In addition to assignment-based type narrowing, Pyright supports the following t * `x is None` and `x is not None` * `x == None` and `x != None` +* `x is ...` and `x is not ...` +* `x == ...` and `x != ...` * `type(x) is T` and `type(x) is not T` * `x is E` and `x is not E` (where E is a literal enum or bool) * `x == L` and `x != L` (where L is an expression that evaluates to a literal type) diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts index dc329cc22..836dd4cac 100644 --- a/packages/pyright-internal/src/analyzer/typeGuards.ts +++ b/packages/pyright-internal/src/analyzer/typeGuards.ts @@ -164,6 +164,22 @@ export function getTypeNarrowingCallback( } } + // Look for "X is ...", "X is not ...", "X == ...", and "X != ...". + if (testExpression.rightExpression.nodeType === ParseNodeType.Ellipsis) { + // Allow the LHS to be either a simple expression or an assignment + // expression that assigns to a simple name. + let leftExpression = testExpression.leftExpression; + if (leftExpression.nodeType === ParseNodeType.AssignmentExpression) { + leftExpression = leftExpression.name; + } + + if (ParseTreeUtils.isMatchingExpression(reference, leftExpression)) { + return (type: Type) => { + return narrowTypeForIsEllipsis(evaluator, type, adjIsPositiveTest); + }; + } + } + // Look for "type(X) is Y" or "type(X) is not Y". if (isOrIsNotOperator && testExpression.leftExpression.nodeType === ParseNodeType.Call) { if ( @@ -876,6 +892,49 @@ function narrowTypeForIsNone(evaluator: TypeEvaluator, type: Type, isPositiveTes ); } +// Handle type narrowing for expressions of the form "x is ..." and "x is not ...". +function narrowTypeForIsEllipsis(evaluator: TypeEvaluator, type: Type, isPositiveTest: boolean) { + const expandedType = mapSubtypes(type, (subtype) => { + return transformPossibleRecursiveTypeAlias(subtype); + }); + + return evaluator.mapSubtypesExpandTypeVars( + expandedType, + /* conditionFilter */ undefined, + (subtype, unexpandedSubtype) => { + if (isAnyOrUnknown(subtype)) { + // We need to assume that "Any" is always both None and not None, + // so it matches regardless of whether the test is positive or negative. + return subtype; + } + + // If this is a TypeVar that isn't constrained, use the unexpanded + // TypeVar. For all other cases (including constrained TypeVars), + // use the expanded subtype. + const adjustedSubtype = + isTypeVar(unexpandedSubtype) && unexpandedSubtype.details.constraints.length === 0 + ? unexpandedSubtype + : subtype; + + // See if it's a match for object. + if (isClassInstance(subtype) && ClassType.isBuiltIn(subtype, 'object')) { + return isPositiveTest + ? addConditionToType(NoneType.createInstance(), subtype.condition) + : adjustedSubtype; + } + + const isEllipsis = isClassInstance(subtype) && ClassType.isBuiltIn(subtype, 'ellipsis'); + + // See if it's a match for "...". + if (isEllipsis === isPositiveTest) { + return subtype; + } + + return undefined; + } + ); +} + // The "isinstance" and "issubclass" calls support two forms - a simple form // that accepts a single class, and a more complex form that accepts a tuple // of classes (including arbitrarily-nested tuples). This method determines diff --git a/packages/pyright-internal/src/tests/samples/typeNarrowingIsEllipsis1.py b/packages/pyright-internal/src/tests/samples/typeNarrowingIsEllipsis1.py new file mode 100644 index 000000000..51b58c38e --- /dev/null +++ b/packages/pyright-internal/src/tests/samples/typeNarrowingIsEllipsis1.py @@ -0,0 +1,37 @@ +# This sample tests the type analyzer's type narrowing logic for +# conditions of the form "X is ...", "X is not ...", +# "X == .." and "X != ...". + +import types +from typing import TypeVar + + +_T = TypeVar("_T", str, ellipsis) + + +def func1(val: int | ellipsis): + if val is not ...: + reveal_type(val, expected_text="int") + else: + reveal_type(val, expected_text="ellipsis") + + +def func2(val: _T): + if val is ...: + reveal_type(val, expected_text="ellipsis*") + else: + reveal_type(val, expected_text="str*") + + +def func3(val: int | types.EllipsisType): + if val != ...: + reveal_type(val, expected_text="int") + else: + reveal_type(val, expected_text="ellipsis") + + +def func4(val: int | ellipsis): + if not val == ...: + reveal_type(val, expected_text="int") + else: + reveal_type(val, expected_text="ellipsis") diff --git a/packages/pyright-internal/src/tests/typeEvaluator1.test.ts b/packages/pyright-internal/src/tests/typeEvaluator1.test.ts index afd84d692..0c5489f62 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator1.test.ts +++ b/packages/pyright-internal/src/tests/typeEvaluator1.test.ts @@ -321,6 +321,12 @@ test('TypeNarrowingIsNoneTuple2', () => { TestUtils.validateResults(analysisResults, 0); }); +test('TypeNarrowingIsEllipsis1', () => { + const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingIsEllipsis1.py']); + + TestUtils.validateResults(analysisResults, 0); +}); + test('TypeNarrowingLiteral1', () => { const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingLiteral1.py']);