Looping through literals not typed correctly

[rggjan]

Iterating through a fixed Tuple of strings ("foo", "bar") makes the loop variable a str instead of Union[Literal["foo"], Literal["bar"]]. This makes it difficult to loop through indices of a TypedDict.´

https://mypy-play.net/?mypy=latest&python=3.8&gist=17fe6a875f727a01fe3a5c6dca13dba2

from typing import TypedDict

class FooDict(TypedDict):
    foo: int
    bar: int
    
foo = FooDict(foo=3, bar=3)

print(foo["foo"]) # Works
print(foo["bar"]) # Works
reveal_type(("foo", "bar")) # Revealed type is 'Tuple[Literal['foo']?, Literal['bar']?]'

for key in ("foo", "bar"):
    reveal_type(key) # Revealed type is 'builtins.str'
    print(foo[key]) # TypedDict key must be a string literal; expected one of ('foo', 'bar')

[Akuli]

is this a duplicate of #9168 ?

edit: no but it's related

[rggjan]

Yes, it's a pretty similar case, I agree...

[JukkaL]

The problem with inferring a literal type here is that then code like this could generate a false positive:

for x in ('foo', 'bar'):
    x = x.upper()  # str is not compatible with a literal type
    print(x)

This would be more feasible if mypy would allow freely redefining variables with different types.

[JukkaL]

Actually, we could maybe infer str as the actual type of x, and narrow it down to a union of literal types in the body of the for loop. I think that this might work.

[rggjan]

I see the potential issue. But currently even this fails:

key: Literal["foo", "bar"]
for key in ("foo", "bar"): # error: Incompatible types in assignment (expression has type "str", variable has type "Union[Literal['foo'], Literal['bar']]")
    print(foo[key]) # TypedDict key must be a string literal; expected one of ('foo', 'bar')

which makes the issue very hard to work around...

[Akuli]

for x in ('foo', 'bar'):
   x = x.upper()  # str is not compatible with a literal type
   print(x)

Why would anyone want to do this? If a variable comes from looping over hard-coded strings, then why would you ever want to change it rather than looping over different hard-coded strings, like this:

for x in ('FOO', 'BAR'):
    print(x)

I guess the only situation is if you want to use both the lowercase x and the uppercase x, but for different things (but why not just create two variables then?)

for x in ('foo', 'bar'):
    print(x)
    x = x.upper()
    print(x)

I guess we should somehow search a big amount of python code to see whether x = x.upper() not supported for Literals is actually a problem.

Similarly, mypy disallows x = x.split(). Have people complained about that?

[gvanrossum]

I have seen this pattern enough times that you needn’t go on a hunt. For example the strings may be keys and the capitalized version will be presented to the user. Etc., etc.

[Akuli]

for x in ('foo', 'bar'):
    print(x)
    x = x.upper()
    print(x)

mypy can already "narrow down" the type of a local variable. For example, if foo has type Any (or e.g. object or Union[int, str]), then assert isinstance(foo, int) changes the type of foo to int.

Maybe there should also be a way to "widen up" the type of a local variable? In this case, x = x.upper() would change the type of x from Literal['foo', 'bar'] to str. Or maybe just support putting x: str before the loop?

This wouldn't be great even if it worked...

key: Literal["foo", "bar"]
for key in ("foo", "bar"):
    # key has type Literal['foo', 'bar']
    ...

...because "foo", "bar" needs to be spelled twice which makes typos possible. But with modifications only in typeshed, I think it might be possible to make this work:

for key in typing_extensions.get_args(Literal['foo', 'bar']):
    # key has type Literal['foo', 'bar']
    ...

Edit: simplified last example code

[Akuli]

that actually won't work with modifications in typeshed only:

def literal_values(lit: Type[T]) -> T:
    return cast(Iterable[T], get_args(lit))

reveal_type(literal_values(Literal['foo', 'bar']))  # <nothing>

[Dr-Irv]

I see the potential issue. But currently even this fails:

key: Literal["foo", "bar"]
for key in ("foo", "bar"): # error: Incompatible types in assignment (expression has type "str", variable has type "Union[Literal['foo'], Literal['bar']]")
    print(foo[key]) # TypedDict key must be a string literal; expected one of ('foo', 'bar')

which makes the issue very hard to work around...

Here is something that worked for me using typing.get_args:

FooBarType = Literal["foo", "bar"]
for key in get_args(FooBarType):
    print(key)

This allows you to loop through all the possible Literal values, although the type of key is not FooBarType, but if you pass it to a function/method expecting FooBarType, mypy does not complain.