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.

[solsword]

Sadly, the workaround listed above results in a type of Any for the loop variable. However, you can add a # type: comment to fix this and avoid the troubles that a stray unintented Any can bring:

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

In the above, mypy does not complain because the type of the get_args item is Any. In this modified version, the type is the Literal you just defined, and mypy can check uses against that.

[Dr-Irv]

Sadly, the workaround listed above results in a type of Any for the loop variable. However, you can add a # type: comment to fix this and avoid the troubles that a stray unintented Any can bring:

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

I was unaware of type comments, and they are soon to be removed, so the appropriate way of doing this would be to write:

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

[jacob-bush-shopify]

I would like to mention that the get_args suggestion is a work around not a solution.
For instance, this passes mypy --strict (my machine is running python==3.11.8, mypy==1.9.0):

from typing import Literal, get_args

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

But when executed, this of course prints "baz"

[jacob-bush-shopify]

Though a bit verbose, here is my suggestion:

from typing import Literal

FooBar = Literal["foo", "bar"]
FOO_BARS: tuple[FooBar, ...] = ("foo", "bar")

for key in FOO_BARS:
    print(key)

[Jeitan]

Just found this thread because I ran into the issue that not even .keys() works - pretty much as expected, but it's a slightly different use-case than above where the strings being looped over are already spelled out somewhere. I've hit it because I want to access my dict by its keys in a loop, something like this a la the OP's setup:

from typing import TypedDict

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

for key in foo:
    print(foo[key])  # TypedDict key must be a string literal; expected one of ("foo", "bar")

While the workaround from @JacobBush does indeed work (ty!), it adds extra bloat and I would think this is a pretty common use-case.

[Dr-Irv]

@Jeitan for what it's worth, pyright does not complain about your example.

[Jeitan]

@Dr-Irv Huh. Thanks for the tip, although it doesn't particularly help on the system of interest. I wonder what is different between the two.

FWIW I got around it by making a Literal type with all the keys, then cast to that inside the loop.