in keyword should narrow Union[TypedDict, ...]

[mashumafi]

Feature

Using in on a Union[TypedDict] should narrow the Union. There is already something similar for isinstance that can narrow based on types. If a certain key exists you could narrow down that your types from a Union as well. This would prevent an error saying a certain key is missing because you checked already, they code can safely access the key. Currently the only workaround for this would be to use cast() but really shouldn't be necessary.

Pitch

After the change the following code should pass validation.

from typing import castm Union
from typing_extensions import TypedDict

class Movie(TypedDict):
    name: str
    year: int

class MovieResponse(TypedDict):
    movie: Movie

class Book(TypedDict):
    name: str
    year: int

class BookResponse(TypedDict):
    book: Book

MediaResponse = Union[MovieResponse, BookResponse]

response : MediaResponse = MovieResponse(movie=Movie(name="Blade Runner", year=1982))

if 'movie' in response:
    reveal_type(response)  # Current: Union[MovieResponse, BookResponse]
                           # Desired: Union[MovieResponse]
    print(response["movie"])  # This fails due to `BookResponse` not having a key "movie"
                              # but we already validated existance of the 'movie' key and that type
                              # should no longer be considered here
    print(cast(MovieResponse, response)["movie"])  # current workaround

[michaeloliverx]

You can get around this by giving each typeddict a unique tag but it would be nice not to have to do this. TypeScript has this feature "The in operator narrowing"

[choucavalier]

indeed this would be great! having mypy errors i should not have because of this

[ikonst]

Since #13838 this already works, but you have to decorate the TypedDicts with @final:

from typing import final

@final
class MovieResponse(TypedDict):
  ...

@final
class BookResponse(TypedDict):
  ...

since otherwise you can't rule out

class BookBasedOnMovieResponse(BookResponse):
  movie: Movie

@AlexWaygood we could close this issue, but I do think we could better document it at the very least (#15695).

p.s. I also wonder if mypy should surface a hint, e.g.

notice: "'foo' in Union[HasFoo, HasBar]" would not exclude HasBar since it's not marked @final"

[choucavalier]

Thanks @ikonst indeed this worked!

[ikonst]

FWIW we're considering removing the need to add @final: #15697