Improved type inference logic for tuple expressions that contain unpa…

…cked tuples.
microsoft · Feb 25, 2021 · b65928c · b65928c
1 parent f5916b7
commit b65928c
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 12 deletions.
diff --git a/packages/pyright-internal/src/analyzer/typeEvaluator.ts b/packages/pyright-internal/src/analyzer/typeEvaluator.ts
@@ -5372,7 +5372,9 @@ export function createTypeEvaluator(
     }
 
     function buildTupleTypesList(entryTypeResults: TypeResult[]): Type[] {
-        let tupleTypes: Type[] = [];
+        const entryTypes: Type[] = [];
+        let isOpenEnded = false;
+
         for (const typeResult of entryTypeResults) {
             if (typeResult.unpackedType) {
                 // Is this an unpacked tuple? If so, we can append the individual
@@ -5384,24 +5386,26 @@ export function createTypeEvaluator(
 
                     // If the Tuple wasn't specialized or has a "..." type parameter, we can't
                     // make any determination about its contents.
-                    if (!typeArgs || typeArgs.some((t) => isEllipsisType(t))) {
-                        tupleTypes = [AnyType.create(/* isEllipsis */ false), AnyType.create(/* isEllipsis */ true)];
-                        break;
-                    }
-
-                    for (const typeArg of typeArgs) {
-                        tupleTypes.push(typeArg);
+                    if (!typeArgs || isOpenEndedTupleClass(typeResult.unpackedType.classType)) {
+                        entryTypes.push(typeResult.type);
+                        isOpenEnded = true;
+                    } else {
+                        entryTypes.push(...typeArgs);
                     }
                 } else {
-                    tupleTypes = [AnyType.create(/* isEllipsis */ false), AnyType.create(/* isEllipsis */ true)];
-                    break;
+                    entryTypes.push(UnknownType.create());
+                    isOpenEnded = true;
                 }
             } else {
-                tupleTypes.push(typeResult.type);
+                entryTypes.push(typeResult.type);
             }
         }
 
-        return tupleTypes;
+        if (isOpenEnded) {
+            return [combineTypes(entryTypes), AnyType.create(/* isEllipsis */ true)];
+        }
+
+        return entryTypes;
     }
 
     function updateNamedTupleBaseClass(classType: ClassType, typeArgs: Type[], isTypeArgumentExplicit: boolean) {

diff --git a/packages/pyright-internal/src/tests/samples/tuples12.py b/packages/pyright-internal/src/tests/samples/tuples12.py
@@ -0,0 +1,19 @@
+# This sample tests type inference for tuples that contain unpack
+# operators.
+
+from typing import Literal
+
+
+def func1(a: int, *args: int):
+    v1 = (a, *args)
+    t1: Literal["tuple[int, ...]"] = reveal_type(v1)
+
+
+def func2(a: int, *args: str):
+    v1 = (a, *args)
+    t1: Literal["tuple[int | str, ...]"] = reveal_type(v1)
+
+
+def func3(a: int, b: str, *args: str):
+    v1 = (*args, a, *args, b, *(a, b, a))
+    t1: Literal["tuple[str | int, ...]"] = reveal_type(v1)
diff --git a/packages/pyright-internal/src/tests/typeEvaluator1.test.ts b/packages/pyright-internal/src/tests/typeEvaluator1.test.ts
@@ -834,6 +834,12 @@ test('Tuples11', () => {
     TestUtils.validateResults(analysisResults, 1);
 });
 
+test('Tuples12', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['tuples12.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('NamedTuples1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['namedTuples1.py']);