Added support for new type guard pattern: x[I] is None and x[I] is not None where x is a tuple or union of tuples with known lengths and entry types and I is an integer.

This commit is contained in:
Eric Traut 2021-10-06 22:48:14 -07:00
parent ff55b895f9
commit 1cbb690b11
4 changed files with 139 additions and 34 deletions

View File

@ -155,6 +155,7 @@ In addition to assignment-based type narrowing, Pyright supports the following t
* `x.y == L` and `x.y != L` (where L is a literal expression and x is a type that is distinguished by a field with a literal type)
* `x[K] == V` and `x[K] != V` (where K and V are literal expressions and x is a type that is distinguished by a TypedDict field with a literal type)
* `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 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)
* `x in y` (where y is instance of list, set, frozenset, or deque)
* `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

@ -117,42 +117,25 @@ export function getTypeNarrowingCallback(
}
if (ParseTreeUtils.isMatchingExpression(reference, leftExpression)) {
// Narrow the type by filtering on "None".
return (type: Type) => {
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;
}
return narrowTypeForIsNone(evaluator, type, adjIsPositiveTest);
};
}
// 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 adjIsPositiveTest ? NoneType.createInstance() : adjustedSubtype;
}
// See if it's a match for None.
if (isNone(subtype) === adjIsPositiveTest) {
return adjustedSubtype;
}
return undefined;
}
);
if (
leftExpression.nodeType === ParseNodeType.Index &&
ParseTreeUtils.isMatchingExpression(reference, leftExpression.baseExpression) &&
leftExpression.items.length === 1 &&
!leftExpression.trailingComma &&
leftExpression.items[0].argumentCategory === ArgumentCategory.Simple &&
!leftExpression.items[0].name &&
leftExpression.items[0].valueExpression.nodeType === ParseNodeType.Number &&
leftExpression.items[0].valueExpression.isInteger &&
!leftExpression.items[0].valueExpression.isImaginary
) {
const indexValue = leftExpression.items[0].valueExpression.value;
return (type: Type) => {
return narrowTupleTypeForIsNone(evaluator, type, adjIsPositiveTest, indexValue);
};
}
}
@ -460,6 +443,79 @@ export function getTypeNarrowingCallback(
return undefined;
}
// Handle type narrowing for expressions of the form "a[I] is None" and "a[I] is not None" where
// I is an integer and a is a union of Tuples with known lengths and entry types.
function narrowTupleTypeForIsNone(evaluator: TypeEvaluator, type: Type, isPositiveTest: boolean, indexValue: number) {
return evaluator.mapSubtypesExpandTypeVars(type, /* conditionFilter */ undefined, (subtype) => {
if (
!isClassInstance(subtype) ||
!isTupleClass(subtype) ||
isOpenEndedTupleClass(subtype) ||
!subtype.tupleTypeArguments
) {
return subtype;
}
const tupleLength = subtype.tupleTypeArguments.length;
if (indexValue < 0 || indexValue >= tupleLength) {
return subtype;
}
const typeOfEntry = evaluator.makeTopLevelTypeVarsConcrete(subtype.tupleTypeArguments[indexValue]);
if (isPositiveTest) {
if (!evaluator.canAssignType(typeOfEntry, NoneType.createInstance(), new DiagnosticAddendum())) {
return undefined;
}
} else {
if (isNone(typeOfEntry)) {
return undefined;
}
}
return subtype;
});
}
// Handle type narrowing for expressions of the form "x is None" and "x is not None".
function narrowTypeForIsNone(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 ? NoneType.createInstance() : adjustedSubtype;
}
// See if it's a match for None.
if (isNone(subtype) === isPositiveTest) {
return adjustedSubtype;
}
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. This method determines which form and returns a list of classes

View File

@ -0,0 +1,42 @@
# This sample tests the type narrowing case for unions of tuples
# where one or more of the entries is tested against type None.
from typing import Literal, TypeVar, Tuple, Union
_T1 = TypeVar("_T1")
def func1(a: Union[Tuple[_T1, None], Tuple[None, str]]) -> Tuple[_T1, None]:
if a[1] is None:
t1: Literal["Tuple[_T1@func1, None]"] = reveal_type(a)
return a
else:
t2: Literal["Tuple[None, str]"] = reveal_type(a)
raise ValueError()
_T2 = TypeVar("_T2", bound=Union[None, int])
def func2(a: Union[Tuple[_T2, None], Tuple[None, str]]):
if a[0] is None:
t1: Literal["Tuple[_T2@func2, None] | Tuple[None, str]"] = reveal_type(a)
else:
t2: Literal["Tuple[_T2@func2, None]"] = reveal_type(a)
_T3 = TypeVar("_T3", None, int)
def func3(a: Union[Tuple[_T3, None], Tuple[None, str]]):
if a[0] is None:
t1: Literal["Tuple[_T3@func3, None] | Tuple[None, str]"] = reveal_type(a)
else:
t2: Literal["Tuple[_T3@func3, None]"] = reveal_type(a)
def func4(a: Union[Tuple[Union[int, None]], Tuple[None, str]]):
if a[0] is None:
t1: Literal["Tuple[int | None] | Tuple[None, str]"] = reveal_type(a)
else:
t2: Literal["Tuple[int | None]"] = reveal_type(a)

View File

@ -295,6 +295,12 @@ test('TypeNarrowingIsNone2', () => {
TestUtils.validateResults(analysisResults, 0);
});
test('TypeNarrowingIsNoneTuple1', () => {
const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingIsNoneTuple1.py']);
TestUtils.validateResults(analysisResults, 0);
});
test('TypeNarrowingLiteral1', () => {
const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingLiteral1.py']);