Stronger type-annotation for a mixed-type mapping

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.

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: ...