diff --git a/packages/pyright-internal/src/analyzer/typeGuards.ts b/packages/pyright-internal/src/analyzer/typeGuards.ts index 9c5e69bc3612..82f2353e0f72 100644 --- a/packages/pyright-internal/src/analyzer/typeGuards.ts +++ b/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. diff --git a/packages/pyright-internal/src/tests/samples/typedDictClosed11.py b/packages/pyright-internal/src/tests/samples/typedDictClosed11.py new file mode 100644 index 000000000000..dfbe6e4a624c --- /dev/null +++ b/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") diff --git a/packages/pyright-internal/src/tests/typeEvaluator5.test.ts b/packages/pyright-internal/src/tests/typeEvaluator5.test.ts index cbed0fdbcb9a..b180c4e11bb6 100644 --- a/packages/pyright-internal/src/tests/typeEvaluator5.test.ts +++ b/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']);