diff --git a/docs/type-concepts-advanced.md b/docs/type-concepts-advanced.md index e7cc47b87..69e072c2d 100644 --- a/docs/type-concepts-advanced.md +++ b/docs/type-concepts-advanced.md @@ -71,7 +71,7 @@ In addition to assignment-based type narrowing, Pyright supports the following t * `x[I] == V` and `x[I] != V` (where I and V are literal expressions and x is a known-length tuple that is distinguished by the index indicated by I) * `x[I] is B` and `x[I] is not B` (where I is a literal expression, B is a `bool` or enum literal, and x is a known-length tuple that is distinguished by the index indicated by I) * `x[I] is None` and `x[I] is not None` (where I is a literal expression and x is a known-length tuple that is distinguished by the index indicated by I) -* `len(x) == L` and `len(x) != L` (where x is tuple and L is a literal integer) +* `len(x) == L` and `len(x) != L` (where x is tuple and L is an expression that evaluates to an int literal type) * `x in y` or `x not in y` (where y is instance of list, set, frozenset, deque, tuple, dict, defaultdict, or OrderedDict) * `S in D` and `S not in D` (where S is a string literal and D is a TypedDict) * `isinstance(x, T)` (where T is a type or a tuple of types) diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts index 9f314f344..500341138 100644 --- a/packages/pyright-internal/src/analyzer/typeGuards.ts +++ b/packages/pyright-internal/src/analyzer/typeGuards.ts @@ -410,9 +410,7 @@ export function getTypeNarrowingCallback( if ( equalsOrNotEqualsOperator && testExpression.leftExpression.nodeType === ParseNodeType.Call && - testExpression.leftExpression.arguments.length === 1 && - testExpression.rightExpression.nodeType === ParseNodeType.Number && - testExpression.rightExpression.isInteger + testExpression.leftExpression.arguments.length === 1 ) { const arg0Expr = testExpression.leftExpression.arguments[0].valueExpression; @@ -424,13 +422,20 @@ export function getTypeNarrowingCallback( const callType = callTypeResult.type; if (isFunction(callType) && callType.details.fullName === 'builtins.len') { - const tupleLength = testExpression.rightExpression.value; + const rightTypeResult = evaluator.getTypeOfExpression(testExpression.rightExpression); + const rightType = rightTypeResult.type; + + if ( + isClassInstance(rightType) && + typeof rightType.literalValue === 'number' && + rightType.literalValue >= 0 + ) { + const tupleLength = rightType.literalValue; - if (typeof tupleLength === 'number') { return (type: Type) => { return { type: narrowTypeForTupleLength(evaluator, type, tupleLength, adjIsPositiveTest), - isIncomplete: !!callTypeResult.isIncomplete, + isIncomplete: !!callTypeResult.isIncomplete || !!rightTypeResult.isIncomplete, }; }; } diff --git a/packages/pyright-internal/src/tests/samples/typeNarrowingTupleLength1.py b/packages/pyright-internal/src/tests/samples/typeNarrowingTupleLength1.py index 0136908c3..922c39d38 100644 --- a/packages/pyright-internal/src/tests/samples/typeNarrowingTupleLength1.py +++ b/packages/pyright-internal/src/tests/samples/typeNarrowingTupleLength1.py @@ -1,6 +1,6 @@ # This sample tests type narrowing of tuples based on len(x) test. -from typing import TypeVar +from typing import Literal, TypeVar def func1(val: tuple[int] | tuple[int, int] | tuple[str, str]): @@ -28,7 +28,8 @@ def func2(val: tuple[int] | tuple[int, ...]): def func3(val: tuple[int] | tuple[()]): - if len(val) == 0: + N = 0 + if len(val) == N: reveal_type(val, expected_text="tuple[()]") else: reveal_type(val, expected_text="tuple[int]") @@ -52,9 +53,10 @@ def func5( | tuple[str] | tuple[str, str, str] | tuple[int, *tuple[str, ...], str] - | tuple[int, *tuple[float, ...]] + | tuple[int, *tuple[float, ...]], + length: Literal[2], ): - if len(val) == 2: + if len(val) == length: reveal_type( val, expected_text="tuple[int, int] | tuple[int, str] | tuple[int, float]" )