Inferring Literal types not always done

[Dr-Irv]

Describe the bug
If you loop over values corresponding to a Literal type, then pyright doesn't infer that the values in the loop are from the Literal type

To Reproduce
See example below. pyright reports issues with the assignment of d1 and d2. Note that mypy reports an issue with the assignment of d1 but not d2

Expected behavior
Ideally, neither line would be flagged. But it would be at least an improvement if pyright didn't report the error on the assignment of d2 that mypy accepts

Screenshots or Code

from typing import Dict, Literal, List

LitChoices = Literal["ab", "bcd"]


keys: List[LitChoices] = ["ab", "bcd"]

d1: Dict[LitChoices, int] = {k: len(k) for k in ["ab", "bcd"]}
d2: Dict[LitChoices, int] = {k: len(k) for k in keys}


print(d1, d2)

VS Code extension or command-line

pyright 1.1.114 command line (also 1.1.112 in VS Code)

Additional context

This kind of code is kind of ugly, because you have to specify the possible values twice - in the type declaration and in the list. Ideally, there would be ONE place to show the possible values.

A possible workaround is to use Enum, as follows, but pyright complains here about the hack (found at python/typing#535 (comment)) in the return of choices (as does mypy):

from enum import Enum
from typing import TypeVar, Tuple, Type, List, Dict

EGen = TypeVar("EGen")


class EnumIt(Enum):
    @classmethod
    def choices(cls: Type[EGen]) -> List[Tuple[EGen, str]]:
        return [(item, name) for name, item in cls.__members__.items()]


class Lets(EnumIt):
    ab = "ab"
    bcd = "bcd"


d3: Dict[Lets, int] = {k: 3 * len(k.name) for k, _ in Lets.choices()}

print(d3)

[erictraut]

The first of your two examples (d1: Dict[LitChoices, int] = {k: len(k) for k in ["ab", "bcd"]}) is expected. That's because of the type inference rules used for list expressions. The inferred type of this expression is List[str], and bidirectional type inference is not supported for for statements.

The second of your two examples (d2: Dict[LitChoices, int] = {k: len(k) for k in keys}) I would expect would work. I'll investigate further.

[erictraut]

The second case in your example (d2) was a bug that will be fixed in the next release. As I mentioned above, the first case (d1) is expected behavior.

[erictraut]

This is addressed in Pyright 1.1.115, which was just published. It will also be in the next release of Pylance.