Narrowing Union[TypedDict] with `in` keyword

[devin-git]

Feature

mypy should narrow Union[TypedDict] when using in keyword.

Pitch

from typing import Literal, TypedDict, Union

Key = Literal['a', 'b']

class A(TypedDict):
    a: int
    
class B(TypedDict):
    b: int

def try_print(x: Union[A, B], key: Key):
    if key in x:
        print(x[key])

The example above triggers errors in mypy:

main.py:13: error: TypedDict "A" has no key "b"
main.py:13: error: TypedDict "B" has no key "a"

In theory, if key in x condition should provide enough information for mypy to narrow the type of parameter x. It should know that x is aligned with key, so the print statement is type safe.

Current workaround is to manually cast the type as below,

    if key == 'a':
        print(cast(A, x)[key])
    else:
        print(cast(B, x)[key])

but when there's much more keys and TypedDicts involved, this approach would become verbose and messy.

If mypy can add this narrowing feature, it will allow more flexibility for TypedDict Union use cases

[erictraut]

The condition if key in x unfortunately doesn't allow for any type narrowing. If key can be either Literal["a"] or Literal["b"], then x can be either A or B — i.e. it's a union of A and B.

It sounds like you want some form of conditional type narrowing where mypy determines "if key is 'a' then the type of x must be A, otherwise the type of x must be B". I'm not aware of any static type analyzer (in any language) that performs this type of conditional type tracking. It would be very complex and computationally expensive to do so.

Also, since your classes A and B are not marked @final, they could be subclassed. That means it's not even safe to narrow in this example:

if 'a' in x:
    reveal_type(x) # A | B (because it's not safe to narrow)

If you mark both A and B as @final, then it's theoretically possible to narrow as follows:

if 'a' in x:
    reveal_type(x) # A

I recently added support for this form of narrowing in pyright, but it doesn't appear to be implemented yet in mypy. There's some discussion of it here.

[devin-git]

Hi Eric, thanks for the quick response and explanation. I didn't know this would be tricky feature.

There was a similar feature request from 2019 and it was implemented in this pr. After reading your explanation, I think what makes my use case complex is key being a variable.

Since it's diffcult to do type narrowing here, I'm thinking about poential workaround. Can mypy check whether x[key] is under key in x condition? If the code already ensures key is in x, It shouldn't throw error that x doesn't contain key.

Also thanks for pointing out the potential subclass issue. I just did some experiment and it lookslike mypy doesn't allow me to apply @final to TypedDict: error: @final cannot be used with TypedDict. Maybe they have made some changes after the "final TypedDict" discussion?

Great to know pyright supports narrowing for literal key checking. I really hope mypy will add the same feature

[erictraut]

There are two ways I can think to implement what you're suggesting. The first is what I mentioned above. Mypy would need to track expression types that are conditioned on the types of other expressions. And it would need to "kill" those conditions if the dependent type changes (e.g. if key were reassigned within the guarded code block). Like I said, I don't know of any static type checker that is able to do this. This would involve significant complexity and runtime overhead.

Another potential approach is for mypy to analyze the entire method multiple times when the annotated parameter types are unions. It would need to do this combinatorially. In your example, it would need to analyze try_print four times:

1. x is A and key is Literal['a']
2. x is A and key is Literal['b']
3. x is B and key is Literal['a']
4. x is B and key is Literal['b']

While this approach would work, it has some major downsides: 1) it would significantly slow down type analysis — exponentially in cases where functions accept many parameters with unions, 2) it would produce multiple redundant error messages, once for each pass, 3) it would be a pretty major overhaul of the type checker logic.

[devin-git]

Thanks a lot for sharing your insights into type analyser. I will leave this issue open for a few days, in case someone else is also interested.

[devin-git]

Closed as it's too difficult to implement