Extended the len(x) == L type guard logic to support arbitrary expressions L that evaluate to a literal int type. This addresses #6216. (#6219)

This commit is contained in:
Eric Traut 2023-10-21 12:09:19 -07:00 committed by GitHub
parent b360d00ada
commit 012dc1c182
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 18 additions and 11 deletions

View File

@ -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)

View File

@ -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,
};
};
}

View File

@ -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]"
)