Narrow closed TypedDicts on a key membership check

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

@@ -2277,11 +2277,13 @@ function narrowTypeForTypedDictKey(
                     //     return subtype;
                     // }
 
-                    // if (isNever(tdEntry.valueType)) {
-                    //     // If the entry is typed as Never or the "extra items" is typed as Never,
-                    //     // then this key cannot be present in the TypedDict, and we can eliminate it.
-                    //     return undefined;
-                    // }
+                    if (isNever(tdEntry.valueType)) {
+                        // If the entry is typed as Never or the "extra items" is typed as Never,
+                        // then this key cannot be present in the TypedDict, and we can eliminate it.
+                        // A closed TypedDict has an "extra items" of Never, so this is what allows
+                        // a key check to discriminate between closed TypedDicts.
+                        return undefined;
+                    }
 
                     // If the entry is currently not required and not marked provided, we can mark
                     // it as provided after this guard expression confirms it is.

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

@@ -0,0 +1,99 @@
+# This sample tests narrowing of a union of closed TypedDicts based on
+# an "in" check for a key. A closed TypedDict cannot contain a key that
+# is not one of its known items, so such a check can discriminate.
+
+
+from typing import Never, NotRequired, TypedDict
+
+
+class Foo(TypedDict, closed=True):
+    foo: int
+
+
+class Bar(TypedDict, closed=True):
+    bar: int
+
+
+def func1(u: Foo | Bar) -> int:
+    if "foo" in u:
+        reveal_type(u, expected_text="Foo")
+        return u["foo"]
+    else:
+        reveal_type(u, expected_text="Bar")
+        return u["bar"]
+
+
+def func2(u: Foo | Bar) -> int:
+    if "bar" not in u:
+        reveal_type(u, expected_text="Foo")
+        return u["foo"]
+    else:
+        reveal_type(u, expected_text="Bar")
+        return u["bar"]
+
+
+class Baz(TypedDict, extra_items=int):
+    baz: int
+
+
+def func3(u: Foo | Baz) -> None:
+    # "Baz" allows extra items, so it cannot be eliminated here.
+    if "foo" in u:
+        reveal_type(u, expected_text="Foo | Baz")
+    else:
+        reveal_type(u, expected_text="Baz")
+
+
+class Open(TypedDict):
+    other: int
+
+
+def func4(u: Foo | Open) -> None:
+    # An open TypedDict without "extra_items" is narrowed on a key check
+    # even though it is not sound to do so; this is idiomatic and is
+    # relied upon in practice.
+    if "foo" in u:
+        reveal_type(u, expected_text="Foo")
+    else:
+        reveal_type(u, expected_text="Open")
+
+
+class NeverItem(TypedDict):
+    always: int
+    never: Never
+
+
+def func5(td: NeverItem) -> None:
+    # A declared item typed as Never can never be present either, so the
+    # same elimination applies to it and not only to the "extra items"
+    # entry synthesized for a closed TypedDict.
+    if "never" in td:
+        reveal_type(td, expected_text="Never")
+    else:
+        reveal_type(td, expected_text="NeverItem")
+
+
+class Left(TypedDict, closed=True):
+    common: int
+    left: int
+
+
+class Right(TypedDict, closed=True):
+    common: int
+    right: NotRequired[int]
+
+
+def func6(u: Left | Right) -> None:
+    # "common" is a required known item of both, so neither is eliminated.
+    if "common" in u:
+        reveal_type(u, expected_text="Left | Right")
+    else:
+        reveal_type(u, expected_text="Never")
+
+
+def func7(td: Right) -> None:
+    # "right" is a known item that is not required, so the subtype is kept
+    # and the key is marked as provided rather than eliminated.
+    if "right" in td:
+        reveal_type(td, expected_text="Right")
+        reveal_type(td["right"], expected_text="int")

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

@@ -408,6 +408,11 @@ test('TypedDictClosed10', () => {
     TestUtils.validateResults(analysisResults, 0);
 });
 
+test('TypedDictClosed11', () => {
+    const analysisResults = TestUtils.typeAnalyzeSampleFiles(['typedDictClosed11.py']);
+    TestUtils.validateResults(analysisResults, 0);
+});
+
 test('DataclassTransform1', () => {
     const analysisResults = TestUtils.typeAnalyzeSampleFiles(['dataclassTransform1.py']);