Suppose I have a mapping that maps string -> int, and int -> string in one dictionary. Is there a way to express that with type annotations in a strict way? Currently I can do Dict[Union[str, int], Union[str, int]], but that allows for str -> str, and int -> int, which I do not actually want.
One might try Union[Dict[str, int], Dict[str, int]], but then only one of these is true for the scope it appears in.
Advertisement
Answer
This can be a solution. If you need, you can re-define other methods.
You can also do a d=cast(mydict,existing_dict).
from typing import overload
class mydict(dict):
@overload
def __getitem__(self, k: int) -> str: ...
@overload
def __getitem__(self, k: str) -> int: ...
def __getitem__(self, k):
return super(mydict, self).__getitem__(k)
@overload
def __setitem__(self, k: int, v: str) -> None: ...
@overload
def __setitem__(self, k: str, v: int) -> None: ...
def __setitem__(self, k, v):
super(mydict, self).__setitem__(k, v)
m = mydict()
m['a'] = 1
m[1] = 'a'
x1: str = m[1]
x2: int = m['a']
m['a'] = 1
m[1] = 'a'
x3: int = m[1] # mypy error
x4: str = m['a'] # mypy error
m[2] = 2 # mypy error
m['b'] = 'b' # mypy error
How @MisterMiyagi suggest, also a Protocol can work:
from typing import overload, Protocol
class mydictprotocol(Protocol):
@overload
def __getitem__(self, k: int) -> str: ...
@overload
def __getitem__(self, k: str) -> int: ...
@overload
def __setitem__(self, k: int, v: str) -> None: ...
@overload
def __setitem__(self, k: str, v: int) -> None: ...