microsoft/pyright
1
1
Fork 0
mirror of https://github.com/microsoft/pyright.git synced 2024-08-16 19:30:26 +03:00
Code Issues Projects Releases Wiki Activity

Improved evaluation logic for tuple slices. It now works in cases where the tuple includes unbounded entries or unpacked TypeVarTuples. This addresses #6906. (#6913)

Browse Source
This commit is contained in:
Eric Traut 2024-01-05 00:13:54 -07:00 committed by GitHub
parent d0100aca93
commit 930ff06e55
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 121 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

106
packages/pyright-internal/src/analyzer/typeEvaluator.ts
View File

@ -7195,7 +7195,7 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
return { type: UnknownType.create() };
}
// Handle the special case where the object is a Tuple and
// Handle the special case where the object is a tuple and
// the index is a constant number (integer) or a slice with integer
// start and end values. In these cases, we can determine
// the exact type by indexing into the tuple type array.
@ -7232,49 +7232,11 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
}
} else if (isClassInstance(valueType) && ClassType.isBuiltIn(valueType, 'slice')) {
const tupleType = getSpecializedTupleType(baseType);
if (tupleType && tupleType.tupleTypeArguments && !isUnboundedTupleClass(tupleType)) {
if (index0Expr.nodeType === ParseNodeType.Slice && !index0Expr.stepValue) {
// Create a local helper function to evaluate the slice parameters.
const getSliceParameter = (expression: ExpressionNode | undefined, defaultValue: number) => {
let value = defaultValue;
if (expression) {
const valType = getTypeOfExpression(expression).type;
if (
isClassInstance(valType) &&
ClassType.isBuiltIn(valType, 'int') &&
isLiteralType(valType) &&
typeof valType.literalValue === 'number'
) {
value = valType.literalValue;
if (value < 0) {
value = tupleType.tupleTypeArguments!.length + value;
}
} else {
value = -1;
}
}
return value;
};
const startValue = getSliceParameter(index0Expr.startValue, 0);
const endValue = getSliceParameter(index0Expr.endValue, tupleType.tupleTypeArguments.length);
if (
startValue >= 0 &&
endValue > 0 &&
endValue <= tupleType.tupleTypeArguments.length &&
tupleClassType &&
isInstantiableClass(tupleClassType)
) {
return {
type: ClassType.cloneAsInstance(
specializeTupleClass(
tupleClassType,
tupleType.tupleTypeArguments.slice(startValue, endValue)
)
),
};
}
if (tupleType && index0Expr.nodeType === ParseNodeType.Slice) {
const slicedTupleType = getSlicedTupleType(tupleType, index0Expr);
if (slicedTupleType) {
return { type: slicedTupleType };
}
}
}
@ -7423,6 +7385,64 @@ export function createTypeEvaluator(importLookup: ImportLookup, evaluatorOptions
};
}
// Given a tuple type and a slice expression, determines the resulting
// type if it can be determined. If not, it returns undefined.
function getSlicedTupleType(tupleType: ClassType, sliceNode: SliceNode): Type | undefined {
// We don't handle step values.
if (sliceNode.stepValue || !tupleType.tupleTypeArguments) {
return undefined;
}
const tupleTypeArgs = tupleType.tupleTypeArguments;
const startValue = getTupleSliceParameter(sliceNode.startValue, 0, tupleTypeArgs);
const endValue = getTupleSliceParameter(sliceNode.endValue, tupleTypeArgs.length, tupleTypeArgs);
if (startValue === undefined || endValue === undefined || endValue < startValue) {
return undefined;
}
const slicedTypeArgs = tupleTypeArgs.slice(startValue, endValue);
return ClassType.cloneAsInstance(specializeTupleClass(tupleType, slicedTypeArgs));
}
function getTupleSliceParameter(
expression: ExpressionNode | undefined,
defaultValue: number,
tupleTypeArgs: TupleTypeArgument[]
): number | undefined {
let value = defaultValue;
if (expression) {
const valType = getTypeOfExpression(expression).type;
if (!isClassInstance(valType) || !ClassType.isBuiltIn(valType, 'int') || !isLiteralType(valType)) {
return undefined;
}
value = valType.literalValue as number;
const unboundedIndex = tupleTypeArgs.findIndex(
(typeArg) => typeArg.isUnbounded || isVariadicTypeVar(typeArg.type)
);
if (value < 0) {
value = tupleTypeArgs.length + value;
if (value < 0) {
return undefined;
} else if (unboundedIndex >= 0 && value <= unboundedIndex) {
return undefined;
}
} else {
if (value >= tupleTypeArgs.length) {
return undefined;
} else if (unboundedIndex >= 0 && value > unboundedIndex) {
return undefined;
}
}
}
return value;
}
function getTypeArgs(node: IndexNode, flags: EvaluatorFlags, options?: GetTypeArgsOptions): TypeResultWithNode[] {
const typeArgs: TypeResultWithNode[] = [];
let adjFlags = flags;

58
packages/pyright-internal/src/tests/samples/tuple13.py
View File

@ -1,6 +1,9 @@
# This sample tests indexing of tuples with slice expressions.
from typing import TypeVarTuple
def func1(val1: tuple[int, str, None], val2: tuple[int, ...]):
x1 = val1[:2]
reveal_type(x1, expected_text="tuple[int, str]")
@ -22,3 +25,58 @@ def func1(val1: tuple[int, str, None], val2: tuple[int, ...]):
x7 = val2[:2]
reveal_type(x7, expected_text="tuple[int, ...]")
def func2(val1: tuple[str, *tuple[int, ...], None]):
x1 = val1[:2]
reveal_type(x1, expected_text="tuple[str | int | None, ...]")
x2 = val1[:1]
reveal_type(x2, expected_text="tuple[str]")
x3 = val1[1:]
reveal_type(x3, expected_text="tuple[*tuple[int, ...], None]")
x4 = val1[1:2]
reveal_type(x4, expected_text="tuple[str | int | None, ...]")
x5 = val1[1:-1]
reveal_type(x5, expected_text="tuple[int, ...]")
x6 = val1[:-1]
reveal_type(x6, expected_text="tuple[str, *tuple[int, ...]]")
x7 = val1[:]
reveal_type(x7, expected_text="tuple[str, *tuple[int, ...], None]")
x8 = val1[2:0]
reveal_type(x8, expected_text="tuple[str | int | None, ...]")
Ts = TypeVarTuple("Ts")
def func3(val1: tuple[str, *Ts, None]):
x1 = val1[:2]
reveal_type(x1, expected_text="tuple[str | Union[*Ts@func3] | None, ...]")
x2 = val1[:1]
reveal_type(x2, expected_text="tuple[str]")
x3 = val1[1:]
reveal_type(x3, expected_text="tuple[*Ts@func3, None]")
x4 = val1[1:2]
reveal_type(x4, expected_text="tuple[str | Union[*Ts@func3] | None, ...]")
x5 = val1[1:-1]
reveal_type(x5, expected_text="tuple[*Ts@func3]")
x6 = val1[:-1]
reveal_type(x6, expected_text="tuple[str, *Ts@func3]")
x7 = val1[:]
reveal_type(x7, expected_text="tuple[str, *Ts@func3, None]")
x8 = val1[2:0]
reveal_type(x8, expected_text="tuple[str | Union[*Ts@func3] | None, ...]")