Precise typing for TypedDict wrappers

[JosuaKrause]

Bug or feature request: feature request (I think)

I tried to use TypedDict for my codebase and ran into some issues. Consider the following code snippet:

from typing import Any, Union, Optional
from typing_extensions import Literal
from mypy_extensions import TypedDict


class Test(TypedDict):
    a: str
    b: int
    c: int


TestKeys = Literal["a", "b", "c"]  # QUESTION: is there a Test.KEYS or something equivalent?
TestValue = Union[str, int]  # QUESTION: is there a Test.VALUE_TYPES or something equivalent?


class Foo:
    def __init__(self):
        self._foo: Optional[Test] = None

    def _set(self) -> None:
        self._foo = {"a": "a", "b": 0, "c": 1}

    def __getitem__(self, key: TestKeys) -> Any:
        self._set()
        assert self._foo is not None
        return self._foo[key]

    def a(self) -> str:
        return self["a"]

    def b(self) -> int:
        return self["b"]


class Bar:
    def __init__(self):
        self._bar: Optional[Test] = None

    def _set(self) -> None:
        self._bar = {"a": "a", "b": 0, "c": 1}

    def __getitem__(self, key: TestKeys) -> TestValue:
        self._set()
        assert self._bar is not None
        return self._bar[key]

    def a(self) -> str:
        return self["a"]  # WRONG: Incompatible return value type (got "Union[str, int]", expected "str")

    def b(self) -> int:
        return self["b"]  # WRONG: Incompatible return value type (got "Union[str, int]", expected "int")


a = Foo().a()
a.find("a")

b = Foo().b()
b.find("a")  # CORRECT: "int" has no attribute "find"

c = Foo()["a"]
c.find("a")

d = Foo()["b"]
d.find("a")  # MISSING: no error reported here (if self._foo is not lazily initialized the error gets correctly reported)

e = Bar().a()
e.find("a")

f = Bar().b()
f.find("a")  # CORRECT: "int" has no attribute "find"

g = Bar()["a"]
g.find("a")  # WRONG: Item "int" of "Union[str, int]" has no attribute "find"

h = Bar()["b"]
h.find("a")  # WRONG: Item "int" of "Union[str, int]" has no attribute "find" (should be "int" has no attribute "find")

(Python 3.8.0, mypy==0.740, mypy-extensions==0.4.3)

I'm not sure what the correct return type signature for a wrapping __getitem__ would be. Both Any and Union seem to erase the actual type coming from the TypedDict and lead to incorrect results. Is there currently a way to achieve the desired behavior here? Also, is there a way to specify Literals for the keys without manually copying them?

Something like

KT = TypeVar("KT", bound=TestKeys)  # or better Test.KEYS

def __getitem__(self, key: KT) -> Test[KT]:
    ...

would be cool if it worked.

[JelleZijlstra]

I think the only current way is to use an overload for every key, which is obviously not great.

[ilevkivskyi]

A feature that would allow this (key types) was discussed a while ago at typing meetup (the primary goal for it was typing for pandas dataframes). We will likely support something like this at some point, but not in near future.

[JosuaKrause]

Workaround that currently works, suggested by @JelleZijlstra , however, it will be cumbersome with more keys.

from typing import Optional, overload
from typing_extensions import Literal
from mypy_extensions import TypedDict


class Test(TypedDict):
    a: str
    b: int
    c: int


class Foo:
    def __init__(self):
        self._foo: Optional[Test] = None

    def _set(self) -> None:
        self._foo = {"a": "a", "b": 0, "c": 1}

    @overload
    def __getitem__(self, key: Literal["a"]) -> str:
        ...

    @overload
    def __getitem__(self, key: Literal["b", "c"]) -> int:
        ...

    def __getitem__(self, key):
        self._set()
        assert self._foo is not None
        return self._foo[key]

    def a(self) -> str:
        return self["a"]

    def b(self) -> int:
        return self["b"]


a = Foo().a()
a.find("a")

b = Foo().b()
b.find("a")  # CORRECT: "int" has no attribute "find"

c = Foo()["a"]
c.find("a")

d = Foo()["b"]
d.find("a")  # CORRECT: "int" has no attribute "find"

Should I close this issue for now?

[ilevkivskyi]

Should I close this issue for now?

No, I think it is useful to have this use case in mind while working on key types.

[A-F-V]

Ideally you would want a way for wrapper[key] to be perfectly substituable for underlying_dict[key] in the eyes of the type system, but that does not seem possible at the moment (unless someone has any solutions?)

[A-F-V]

What I have found is that the type given is usually the union across all possible keys even if it is a literal known before running. I wonder if a change could be made to refine this type if it is known. Like C++ constexpr