Added support for new type narrowing pattern: len(x) == L and len(x) != L where x is a tuple or union of tuples and L is a literal integer value.

docs/type-concepts.md

@@ -156,6 +156,7 @@ In addition to assignment-based type narrowing, Pyright supports the following t
 * `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)
+* `len(x) == L` and `len(x) != L` (where x is tuple and L is a literal integer)
 * `x in y` (where y is instance of list, set, frozenset, deque, or tuple)
 * `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/typeGuards.ts

@@ -297,6 +297,33 @@ export function getTypeNarrowingCallback(
                     }
                 }
             }
+
+            // Look for len(x) == <literal> or len(x) != <literal>
+            if (
+                equalsOrNotEqualsOperator &&
+                testExpression.leftExpression.nodeType === ParseNodeType.Call &&
+                testExpression.leftExpression.arguments.length === 1 &&
+                testExpression.rightExpression.nodeType === ParseNodeType.Number &&
+                testExpression.rightExpression.isInteger
+            ) {
+                const arg0Expr = testExpression.leftExpression.arguments[0].valueExpression;
+
+                if (ParseTreeUtils.isMatchingExpression(reference, arg0Expr)) {
+                    const callType = evaluator.getTypeOfExpression(
+                        testExpression.leftExpression.leftExpression,
+                        /* expectedType */ undefined,
+                        EvaluatorFlags.DoNotSpecialize
+                    ).type;
+
+                    if (isFunction(callType) && callType.details.fullName === 'builtins.len') {
+                        const tupleLength = testExpression.rightExpression.value;
+
+                        return (type: Type) => {
+                            return narrowTypeForTupleLength(evaluator, type, tupleLength, adjIsPositiveTest);
+                        };
+                    }
+                }
+            }
         }
 
         if (testExpression.operator === OperatorType.In) {
@@ -1101,6 +1128,31 @@ function narrowTypeForIsInstance(
     return filteredType;
 }
 
+// Attempts to narrow a union of tuples based on their known length.
+function narrowTypeForTupleLength(
+    evaluator: TypeEvaluator,
+    referenceType: Type,
+    lengthValue: number,
+    isPositiveTest: boolean
+) {
+    return mapSubtypes(referenceType, (subtype) => {
+        const concreteSubtype = evaluator.makeTopLevelTypeVarsConcrete(subtype);
+
+        // If it's not a tuple, we can't narrow it.
+        if (
+            !isClassInstance(concreteSubtype) ||
+            !isTupleClass(concreteSubtype) ||
+            isOpenEndedTupleClass(concreteSubtype) ||
+            !concreteSubtype.tupleTypeArguments
+        ) {
+            return subtype;
+        }
+
+        const tupleLengthMatches = concreteSubtype.tupleTypeArguments.length === lengthValue;
+        return tupleLengthMatches === isPositiveTest ? subtype : undefined;
+    });
+}
+
 // Attempts to narrow a type (make it more constrained) based on an "in" or
 // "not in" binary expression.
 function narrowTypeForContains(evaluator: TypeEvaluator, referenceType: Type, containerType: Type) {

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

@@ -0,0 +1,47 @@
+# This sample tests type narrowing of tuples based on len(x) test.
+
+from typing import Literal, Tuple, TypeVar, Union
+
+
+def func1(val: Union[Tuple[int], Tuple[int, int], Tuple[str, str]]):
+    if len(val) == 1:
+        t1: Literal["Tuple[int]"] = reveal_type(val)
+    else:
+        t2: Literal["Tuple[int, int] | Tuple[str, str]"] = reveal_type(val)
+
+    if len(val) != 2:
+        t3: Literal["Tuple[int]"] = reveal_type(val)
+    else:
+        t4: Literal["Tuple[int, int] | Tuple[str, str]"] = reveal_type(val)
+
+
+def func2(val: Union[Tuple[int], Tuple[int, ...]]):
+    if len(val) == 1:
+        t1: Literal["Tuple[int] | Tuple[int, ...]"] = reveal_type(val)
+    else:
+        t2: Literal["Tuple[int, ...]"] = reveal_type(val)
+
+    if len(val) != 2:
+        t3: Literal["Tuple[int] | Tuple[int, ...]"] = reveal_type(val)
+    else:
+        t4: Literal["Tuple[int, ...]"] = reveal_type(val)
+
+
+def func3(val: Union[Tuple[int], Tuple[()]]):
+    if len(val) == 0:
+        t1: Literal["Tuple[()]"] = reveal_type(val)
+    else:
+        t2: Literal["Tuple[int]"] = reveal_type(val)
+
+
+_T1 = TypeVar("_T1", bound=Tuple[int])
+_T2 = TypeVar("_T2", bound=Tuple[str, str])
+
+
+def func4(val: Union[_T1, _T2]) -> Union[_T1, _T2]:
+    if len(val) == 1:
+        t1: Literal["_T1@func4"] = reveal_type(val)
+    else:
+        t2: Literal["_T2@func4"] = reveal_type(val)
+
+    return val

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

@@ -361,6 +361,12 @@ test('TypeNarrowingIsinstance7', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('TypeNarrowingTupleLength1', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingTupleLength1.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('TypeNarrowingIn1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowingIn1.py']);