CODEWITHSUNDEEP
pythonpython-3.xdictionarytype-hinting
znarf94
znarf94

Reputation: 61

Type hinting dict with key/value type pairs

Is it possible with Python type hints to specify the types of a dictionary's keys and values as pairs ?

For instance :

If I write :

Dict[Union[int, str], Union[int, str]]

it allows str -> str and int -> int, which are not allowed.

And with :

Union[Dict[int, str], Dict[str, int]]

the dictionary can be either a Dict[int, str] or Dict[str, int], but not both at the same time ...

I also looked into TypedDict, but it requires to give explicitly all the keys.

Upvotes: 6

Views: 7115

Answers (1)

LyleCheatham
LyleCheatham

Reputation: 31

If using typing.cast is acceptable for your application, then this can be done by making a class that subclasses Dict that has overrides for __setitem__ and __getitem__, casting your dict to that type. From then on, type checkers will infer correct KeyType: ValueType pairs.

The caveats of this approach are that you can't use it to type check proper dict construction, as that happens before the cast. Additionally, you would need to add more overrides for things like update, __iter__ and possibly other dict methods to get type checking for things beyond __setitem__/__getitem__ dict access.

Example:

from typing import Dict, overload, cast


class KVTypePairedDict(Dict):
    @overload
    def __getitem__(self, key: str) -> int: ...
    @overload
    def __getitem__(self, key: int) -> str: ...
    def __getitem__(self, key): ...

    @overload
    def __setitem__(self, key: str, value: int) -> None: ...
    @overload
    def __setitem__(self, key: int, value: str) -> None: ...
    def __setitem__(self, key, value): ...


test: KVTypePairedDict = cast(KVTypePairedDict, {"foo": 0, 1: "bar"})

# str keys
a: int = test["foo"]
test["foo"] = 0
c: str = test["foo"]  # <-- mypy and PyCharm flag this
test["foo"] = "bar"  # <-- mypy and PyCharm flag this

# int keys
d: str = test[1]
test[1] = "bar"
b: int = test[0]  # <-- mypy and PyCharm flag this
test[1] = 0  # <-- mypy and PyCharm flag this

Upvotes: 3

Related Questions