Added support for "tagged union" type narrowing when the conditional …

…expression is of the form `x[K] == V` or `x[K] != V` where `x` is a union of TypedDict objects and `K` is a literal str key value that refers to a field with a literal type and `V` is a literal value.

docs/type-concepts.md

@@ -153,13 +153,13 @@ In addition to assignment-based type narrowing, Pyright supports the following t
 * `x is E` and `x is not E` (where E is an enum value)
 * `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 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)
 * `issubclass(x, T)` (where T is a type or a tuple of types)
 * `callable(x)`
 * `f(x)` (where f is a user-defined type guard as defined in [PEP 647](https://www.python.org/dev/peps/pep-0647/))
-
 * `x` (where x is any expression that is statically verifiable to be truthy or falsy in all cases)
 
 Expressions supported for type guards include simple names, member access chains (e.g. `a.b.c.d`), the unary `not` operator, the binary `and` and `or` operators, subscripts that are constant numbers (e.g. `a[2]`), and call expressions. Other operators (such as arithmetic operators or other subscripts) are not supported.

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

@@ -17208,6 +17208,37 @@ export function createTypeEvaluator(
                             };
                         }
                     }
+
+                    // Look for X[<literal>] == <literal> or X[<literal>] != <literal>
+                    if (
+                        testExpression.leftExpression.nodeType === ParseNodeType.Index &&
+                        testExpression.leftExpression.items.length === 1 &&
+                        !testExpression.leftExpression.trailingComma &&
+                        testExpression.leftExpression.items[0].argumentCategory === ArgumentCategory.Simple &&
+                        ParseTreeUtils.isMatchingExpression(reference, testExpression.leftExpression.baseExpression)
+                    ) {
+                        const indexType = getTypeOfExpression(
+                            testExpression.leftExpression.items[0].valueExpression
+                        ).type;
+
+                        if (
+                            isObject(indexType) &&
+                            ClassType.isBuiltIn(indexType.classType, 'str') &&
+                            isLiteralType(indexType)
+                        ) {
+                            const rightType = getTypeOfExpression(testExpression.rightExpression).type;
+                            if (isObject(rightType) && rightType.classType.literalValue !== undefined) {
+                                return (type: Type) => {
+                                    return narrowTypeForDiscriminatedDictEntryComparison(
+                                        type,
+                                        indexType,
+                                        rightType,
+                                        adjIsPositiveTest
+                                    );
+                                };
+                            }
+                        }
+                    }
                 }
             }
 
@@ -17742,8 +17773,44 @@ export function createTypeEvaluator(
         return narrowedType;
     }
 
+    // Attempts to narrow a TypedDict type based on a comparison (equal or not
+    // equal) between a discriminating entry type that has a declared literal
+    // type to a literal value.
+    function narrowTypeForDiscriminatedDictEntryComparison(
+        referenceType: Type,
+        indexLiteralType: ObjectType,
+        literalType: ObjectType,
+        isPositiveTest: boolean
+    ): Type {
+        let canNarrow = true;
+
+        const narrowedType = mapSubtypes(referenceType, (subtype) => {
+            if (isObject(subtype) && ClassType.isTypedDictClass(subtype.classType)) {
+                const symbolMap = getTypedDictMembersForClass(subtype.classType);
+                const tdEntry = symbolMap.get(indexLiteralType.classType.literalValue as string);
+
+                if (tdEntry && isLiteralTypeOrUnion(tdEntry.valueType)) {
+                    if (isPositiveTest) {
+                        return canAssignType(tdEntry.valueType, literalType, new DiagnosticAddendum())
+                            ? subtype
+                            : undefined;
+                    } else {
+                        return canAssignType(literalType, tdEntry.valueType, 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 node that has a declared literal type to a
+    // between a discriminating field that has a declared literal type to a
     // literal value.
     function narrowTypeForDiscriminatedFieldComparison(
         referenceType: Type,

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

@@ -0,0 +1,36 @@
+# This sample tests type narrowing based on key accesses
+# to unions of TypedDicts that have fields with literal types.
+
+from typing import Literal, TypedDict, Union
+
+
+class NewJobEvent(TypedDict):
+    tag: Literal["new-job"]
+    job_name: str
+    config_file_path: str
+
+
+class CancelJobEvent(TypedDict):
+    tag: Literal[2]
+    job_id: int
+
+
+class OtherEvent(TypedDict):
+    tag: Literal["other-job"]
+    message: str
+
+
+Event = Union[NewJobEvent, CancelJobEvent, OtherEvent]
+
+
+def process_event(event: Event) -> None:
+    if event["tag"] == "new-job":
+        t1: Literal["NewJobEvent"] = reveal_type(event)
+        event["job_name"]
+    elif event["tag"] == 2:
+        t2: Literal["CancelJobEvent"] = reveal_type(event)
+        event["job_id"]
+    else:
+        t3: Literal["OtherEvent"] = reveal_type(event)
+        event["message"]
+

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

@@ -346,6 +346,12 @@ test('TypeNarrowing21', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('TypeNarrowing22', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typeNarrowing22.py']);
+
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('ReturnTypes1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['returnTypes1.py']);