[feature] Discriminated Union (narrow typed dicts based on field type)

[kbumsik]

Describe the Bug

Pyrefly currently doesn't seem to support Discriminated Unions / Tagged Unions

• mypy doc
• pydantic doc

This feature is widely used, including OpneAI SDK (or output of codegen tools like Stainless)

Here is the simple example from mypy (the same as sandbox link).
The below code passes in mypy, pyright, and etc.

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["cancel-job"]
    job_id: int

Event = Union[NewJobEvent, CancelJobEvent]

def process_event(event: Event) -> None:
    # Since we made sure both TypedDicts have a key named 'tag', it's
    # safe to do 'event["tag"]'. This expression normally has the type
    # Literal["new-job", "cancel-job"], but the check below will narrow
    # the type to either Literal["new-job"] or Literal["cancel-job"].
    #
    # This in turns narrows the type of 'event' to either NewJobEvent
    # or CancelJobEvent.
    if event["tag"] == "new-job":
        print(event["job_name"])
    else:
        print(event["job_id"])

We got error like:

ERROR 27:21-31: TypedDict `CancelJobEvent` does not have key `job_name` [[typed-dict-key-error](https://pyrefly.org/en/docs/error-kinds/#typed-dict-key-error)]
ERROR 29:21-29: TypedDict `NewJobEvent` does not have key `job_id` [[typed-dict-key-error](https://pyrefly.org/en/docs/error-kinds/#typed-dict-key-error)]

Sandbox Link

https://pyrefly.org/sandbox/?code=GYJw9gtgBALgngBwJYDsDmUkQWEMoAySMApiAIYA2ANFACqIkAmAIkgMYy0CqKSYKAFDD2lcgGdxUAHIkA7gCkwAIwCiANxIoYACgYJmbTgEoAXIKiXY5NKcLEyVANoAiFPIC0AKxUuAuhZWPsoA+ijkECR24jAggZbsAsBIaCHJlCQhCOQwABbRscKCohJSAMLkKOwklEpqmtp6jKwcMGbx1rb2pBSUruyV1ZTevgFWUMEhSEx2qDBFGlr4ALxQvPwoTrKKKovatBVVNXV7MAHCTCTAUAjg1ZIhJA26T0t2p8ZQHgB8MgJRHQAxFAAMqoapQOQkKAQciXKDiACuIGhyjAeXozSMMCkuXImig5CgAGsSHAoOFIkwoAByGA2Gm0Yg08RAhHkYDQmBgKBMHk017aVz0tD+GkAOnouSQUhIAA9biRJBsKbhYZRKOS8VI8lzGGyiD1nG5PMEXLQXAMjsMzX5aMpEfhdVB2LkSOxiVBlDUwHJIUgNRTyCBwHI2c74AZYDySMQ3SBuo4+ia5CNlP4oLhE71+oMamn-OKgWy6NKpKhYMiUFJwiHfTq3bBGJnrgLnjTo1BY7qE9sTs82VnDkN+0si+MkNdBTBhTYM8tVimC+ZxuNbnMdNPXJNKSR-MYOjVxADV1Z140ty5JtN94IgA

(Only applicable for extension issues) IDE Information

No response

[yangdanny97]

Thanks for bringing this up, it sounds like a useful case to support for type narrowing.

[kbumsik]

Thanks for bringing this up, it sounds like a useful case to support for type narrowing.

[feature] Discriminated Union (narrow typed dicts based on field type)

Thanks for the quick reply :) Just for clarifying, this can be applied not only to typed dicts but also to regular objects (including dataclass, pydantic objects), tuples, or etc.