Repository navigation
Narrowing Union[TypedDict] with in keyword #11080
Description
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
Activity
erictraut commentedon Sep 9, 2021
The condition
if key in xunfortunately doesn't allow for any type narrowing. Ifkeycan be eitherLiteral["a"]orLiteral["b"], thenxcan be eitherAorB— i.e. it's a union ofAandB.It sounds like you want some form of conditional type narrowing where mypy determines "if
keyis 'a' then the type ofxmust beA, otherwise the type ofxmust beB". 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
AandBare not marked@final, they could be subclassed. That means it's not even safe to narrow in this example:If you mark both
AandBas@final, then it's theoretically possible to narrow as follows: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 commentedon Sep 9, 2021
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
keybeing a variable.Since it's diffcult to do type narrowing here, I'm thinking about poential workaround. Can mypy check whether
x[key]is underkey in xcondition? If the code already ensureskeyis inx, It shouldn't throw error thatxdoesn't containkey.Also thanks for pointing out the potential subclass issue. I just did some experiment and it lookslike mypy doesn't allow me to apply
@finalto 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 commentedon Sep 9, 2021
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
keywere 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_printfour times:xisAandkeyisLiteral['a']xisAandkeyisLiteral['b']xisBandkeyisLiteral['a']xisBandkeyisLiteral['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 commentedon Sep 10, 2021
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 commentedon Sep 12, 2021
Closed as it's too difficult to implement