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

[dchevell]

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).

MyPy also sees this as an error even when defining the literal value inline if my variable is annotated with multiple possible values:

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

Pitch
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.