Added support for new type narrowing pattern for discriminating among tuples. The pattern is V[I] == L or V[I] != L where I is an integer literal, L is another literal value, V is a tuple with a known length and a type at index I that is declared as a literal type.

docs/type-concepts.md

@@ -154,6 +154,7 @@ In addition to assignment-based type narrowing, Pyright supports the following t
 * `x == L` and `x != L` (where L is a literal expression)
 * `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 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)

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

@@ -18008,21 +18008,31 @@ export function createTypeEvaluator(
                             testExpression.leftExpression.items[0].valueExpression
                         ).type;
 
-                        if (
+                        if (isClassInstance(indexType) && isLiteralType(indexType)) {
-                            isClassInstance(indexType) &&
+                            if (ClassType.isBuiltIn(indexType, 'str')) {
-                            ClassType.isBuiltIn(indexType, 'str') &&
+                                const rightType = getTypeOfExpression(testExpression.rightExpression).type;
-                            isLiteralType(indexType)
+                                if (isClassInstance(rightType) && rightType.literalValue !== undefined) {
-                        ) {
+                                    return (type: Type) => {
-                            const rightType = getTypeOfExpression(testExpression.rightExpression).type;
+                                        return narrowTypeForDiscriminatedDictEntryComparison(
-                            if (isClassInstance(rightType) && rightType.literalValue !== undefined) {
+                                            type,
-                                return (type: Type) => {
+                                            indexType,
-                                    return narrowTypeForDiscriminatedDictEntryComparison(
+                                            rightType,
-                                        type,
+                                            adjIsPositiveTest
-                                        indexType,
+                                        );
-                                        rightType,
+                                    };
-                                        adjIsPositiveTest
+                                }
-                                    );
+                            } else if (ClassType.isBuiltIn(indexType, 'int')) {
-                                };
+                                const rightType = getTypeOfExpression(testExpression.rightExpression).type;
+                                if (isClassInstance(rightType) && rightType.literalValue !== undefined) {
+                                    return (type: Type) => {
+                                        return narrowTypeForDiscriminatedTupleComparison(
+                                            type,
+                                            indexType,
+                                            rightType,
+                                            adjIsPositiveTest
+                                        );
+                                    };
+                                }
                             }
                         }
                     }
@@ -18677,6 +18687,40 @@ export function createTypeEvaluator(
         return canNarrow ? narrowedType : referenceType;
     }
 
+    function narrowTypeForDiscriminatedTupleComparison(
+        referenceType: Type,
+        indexLiteralType: ClassType,
+        literalType: ClassType,
+        isPositiveTest: boolean
+    ): Type {
+        let canNarrow = true;
+
+        const narrowedType = mapSubtypes(referenceType, (subtype) => {
+            if (isClassInstance(subtype) && ClassType.isTupleClass(subtype) && !isOpenEndedTupleClass(subtype)) {
+                const indexValue = indexLiteralType.literalValue as number;
+                if (subtype.tupleTypeArguments && indexValue >= 0 && indexValue < subtype.tupleTypeArguments.length) {
+                    const tupleEntryType = subtype.tupleTypeArguments[indexValue];
+                    if (tupleEntryType && isLiteralTypeOrUnion(tupleEntryType)) {
+                        if (isPositiveTest) {
+                            return canAssignType(tupleEntryType, literalType, new DiagnosticAddendum())
+                                ? subtype
+                                : undefined;
+                        } else {
+                            return canAssignType(literalType, tupleEntryType, new DiagnosticAddendum())
+                                ? undefined
+                                : subtype;
+                        }
+                    }
+                }
+            }
+
+            canNarrow = false;
+            return subtype;
+        });
+
+        return canNarrow ? narrowedType : referenceType;
+    }
+
     // Attempts to narrow a type based on a comparison (equal or not equal)
     // between a discriminating field that has a declared literal type to a
     // literal value.

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

@@ -0,0 +1,23 @@
+# This sample tests the type narrowing for known-length tuples
+# that have an entry with a declared literal type.
+
+from typing import Tuple, Union, Literal
+
+MsgA = Tuple[Literal[1], str]
+MsgB = Tuple[Literal[2], float]
+
+Msg = Union[MsgA, MsgB]
+
+
+def func1(m: Msg):
+    if m[0] == 1:
+        t1: Literal["Tuple[Literal[1], str]"] = reveal_type(m)
+    else:
+        t2: Literal["Tuple[Literal[2], float]"] = reveal_type(m)
+
+
+def func2(m: Msg):
+    if m[0] != 1:
+        t1: Literal["Tuple[Literal[2], float]"] = reveal_type(m)
+    else:
+        t2: Literal["Tuple[Literal[1], str]"] = reveal_type(m)

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

@@ -346,6 +346,12 @@ test('TypeNarrowingLiteralMember1', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('TypeNarrowingTuple1', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingTuple1.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('TypeNarrowingTypedDict1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingTypedDict1.py']);