Fixed several bugs related to subscripts with unpack operators, notably when the unpack targets a tuple with a known length. This addresses #10723. (#10971)

Author: erictraut

packages/pyright-internal/src/analyzer/typeEvaluator.ts

@@ -8135,24 +8135,91 @@ export function createTypeEvaluator(
         } else {
             // Package up all of the positionals into a tuple.
             const tupleTypeArgs: TupleTypeArg[] = [];
-            positionalArgs.forEach((arg) => {
-                const typeResult = getTypeOfExpression(arg.d.valueExpr);
-                tupleTypeArgs.push({ type: typeResult.type, isUnbounded: false });
-                if (typeResult.isIncomplete) {
-                    isPositionalIndexTypeIncomplete = true;
+
+            const getDeterministicTupleEntries = (type: Type): TupleTypeArg[] | undefined => {
+                let aggregatedArgs: TupleTypeArg[] | undefined;
+                let isDeterministic = true;
+
+                doForEachSubtype(type, (subtype) => {
+                    if (!isDeterministic) {
+                        return;
+                    }
+
+                    const tupleType = getSpecializedTupleType(subtype);
+                    const tupleTypeArgs = tupleType?.priv.tupleTypeArgs;
+
+                    if (
+                        !tupleTypeArgs ||
+                        tupleTypeArgs.some((entry) => entry.isUnbounded || isTypeVarTuple(entry.type))
+                    ) {
+                        isDeterministic = false;
+                        return;
+                    }
+
+                    if (!aggregatedArgs) {
+                        aggregatedArgs = tupleTypeArgs.map((entry) => ({ type: entry.type, isUnbounded: false }));
+                        return;
+                    }
+
+                    if (aggregatedArgs.length !== tupleTypeArgs.length) {
+                        isDeterministic = false;
+                        return;
+                    }
+
+                    for (let i = 0; i < aggregatedArgs.length; i++) {
+                        aggregatedArgs[i] = {
+                            type: combineTypes([aggregatedArgs[i].type, tupleTypeArgs[i].type]),
+                            isUnbounded: false,
+                        };
+                    }
+                });
+
+                if (!isDeterministic || !aggregatedArgs) {
+                    return undefined;
                 }
-            });
 
-            unpackedListArgs.forEach((arg) => {
-                const typeResult = getTypeOfExpression(arg.d.valueExpr);
-                if (typeResult.isIncomplete) {
-                    isPositionalIndexTypeIncomplete = true;
+                return aggregatedArgs;
+            };
+
+            node.d.items.forEach((arg) => {
+                if (arg.d.argCategory === ArgCategory.Simple) {
+                    const typeResult = getTypeOfExpression(arg.d.valueExpr);
+                    tupleTypeArgs.push({ type: typeResult.type, isUnbounded: false });
+                    if (typeResult.isIncomplete) {
+                        isPositionalIndexTypeIncomplete = true;
+                    }
+                    return;
+                }
+
+                if (arg.d.argCategory === ArgCategory.UnpackedList) {
+                    const typeResult = getTypeOfExpression(arg.d.valueExpr);
+                    if (typeResult.isIncomplete) {
+                        isPositionalIndexTypeIncomplete = true;
+                    }
+
+                    const deterministicEntries = getDeterministicTupleEntries(typeResult.type);
+                    if (deterministicEntries) {
+                        appendArray(tupleTypeArgs, deterministicEntries);
+                        return;
+                    }
+
+                    const iterableType =
+                        getTypeOfIterator(typeResult, /* isAsync */ false, arg.d.valueExpr)?.type ??
+                        UnknownType.create();
+                    tupleTypeArgs.push({ type: iterableType, isUnbounded: true });
                 }
-                const iterableType =
-                    getTypeOfIterator(typeResult, /* isAsync */ false, arg.d.valueExpr)?.type ?? UnknownType.create();
-                tupleTypeArgs.push({ type: iterableType, isUnbounded: true });
             });
 
+            const unboundedCount = tupleTypeArgs.filter((typeArg) => typeArg.isUnbounded).length;
+            if (unboundedCount > 1) {
+                const firstUnboundedIndex = tupleTypeArgs.findIndex((typeArg) => typeArg.isUnbounded);
+                const removedEntries = tupleTypeArgs.splice(firstUnboundedIndex);
+                tupleTypeArgs.push({
+                    type: combineTypes(removedEntries.map((entry) => entry.type)),
+                    isUnbounded: true,
+                });
+            }
+
             positionalIndexType = makeTupleObject(evaluatorInterface, tupleTypeArgs);
         }
 

packages/pyright-internal/src/tests/samples/subscript4.py

@@ -0,0 +1,36 @@
+# This sample tests the handling of unpack operators within
+# a subscript.
+
+from typing import NamedTuple
+
+
+class Recorder[T]:
+    def __getitem__(self, item: T) -> T:
+        return item
+
+
+class OneInt(NamedTuple):
+    value: int
+
+
+class IntStrPair(NamedTuple):
+    first: int
+    second: str
+
+
+recorder_pair: Recorder[tuple[int, str]] = Recorder()
+pair = IntStrPair(1, "value")
+result1 = recorder_pair[*pair]
+reveal_type(result1, expected_text="tuple[int, str]")
+
+recorder_order: Recorder[tuple[int, str]] = Recorder()
+tail_value: str = "tail"
+result2 = recorder_order[*OneInt(2), tail_value]
+reveal_type(result2, expected_text="tuple[int, str]")
+
+recorder_multi: Recorder[tuple[int, *tuple[int | str, ...]]] = Recorder()
+values1: list[int] = []
+values2: list[str] = []
+first_value: int = 0
+result3 = recorder_multi[first_value, *values1, *values2]
+reveal_type(result3, expected_text="tuple[int, *tuple[int | str, ...]]")

packages/pyright-internal/src/tests/typeEvaluator6.test.ts

@@ -964,6 +964,11 @@ test('Subscript3', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('Subscript4', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['subscript4.py']);
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('Decorator1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['decorator1.py']);