Apply TypedDict-like type hint to __getitem__

Python Help
,
1

Can I somehow apply TypedDict-like type hints to __getitem__ method in class to make different type hints for return value based on provided key value?

2

I think you already can, you need to write overloads where each overload key is a typing.Literal string key then the return type of that overload matches the associated key. I think last time I tried it pyright worked as expected, I don’t think I tried with mypy. Something like:

from typing import Literal, overload, assert_type

class A:
    @overload
    def __getitem__(self, key: Literal['foo']) -> int: ...
    @overload
    def __getitem__(self, key: Literal['bar']) -> str: ...


a = A()
assert_type(a['foo'], int)
assert_type(a['bar'], str)
1 Like
4

Thanks for your answer, but can TypedDict be applied to the class (it has less complex syntax than overloads)?

1 Like
5

Unfortunately no, TypedDict is a way to declare the structure of dictionaries, it can’t be used to declare full classes. At runtime I believe types declared using TypedDict are actually just dictionaries and there’s some interpreter magic.

6

Another option would be to look at python dataclasses as these have similar code generation for classes rather than dictionaries. For example:

from dataclasses import dataclass

@dataclass
class Foo:
    bar: int
    baz: str

foo = Foo(1, '2')
print(foo.bar)

Behaves as expected, the dataclass decorator generates you a constructor with the appropriate signature, it also generates a few other methods for you like a nice string representation.

2 Likes