Feature request: infer pre-narrowed tagged union as a Union of possible types

[dchevell]

If you have a question about a behavior that you’re seeing in pyright, consider posting to the Pyright discussion forum.

Is your feature request related to a problem? Please describe.
I'm building a tagged union inside of a function, like so:

from typing import Literal, TypedDict

class Foo(TypedDict):
    tag: Literal["foo"]

class Bar(TypedDict):
    tag: Literal["bar"]
    
def test1(tag: Literal["foo", "bar"]) -> Foo | Bar:
    return {"tag": tag}  # error

Because tag hasn't been narrowed to "foo" or "bar" at the return site, it's considered not to be assignable to FooOrBar, even though there are no other possible outcomes. To fix this, I have to return the same value from both branches of an if clause:

def test2(tag: Literal["foo", "bar"]) -> Foo | Bar:
    if tag == "foo":
        return {"tag": tag}  # ok
    else:
        return {"tag": tag}  # ok

At this stage all I care about is knowing that my output value is one of Foo or Bar, and the above construct looks like a code smell (especially with >2 tag values).

Interestingly, pyright doesn't see this as ambiguous in other contexts:

tag: Literal["foo", "bar"] = "foo"
foo: Foo | Bar = {"tag": tag}  # ok

Describe the solution you’d like
In cases where the "tag" value of the tagged union is not narrowed, but can be inferred as consistent with multiple typed dicts in the return type of a function, I'd love to see this be narrowed as a union of those possible typed dicts. Given the "tag" is already scoped to the possible values (if it weren't, the bare "else" wouldn't work) I can't think of any way this would be unsound.