How to annotate a dict of type to function that returns type

I want to use a dict which maps a type to a callable which returns the same type. E.g.

MY_DICT = {
    bytes: lambda x: x,
    int: lambda x: int.from_bytes(x, 'big'),
    str: lambda x: x.decode('utf-8'),
}

i.e. each value is a callable which takes a bytes value and returns the corresponding type for its key.

I'm wondering how I should annotate this dict for mypy? So far I have come up with:

dict[type, typing.Callable[[bytes], typing.Any]]

But I'm wondering if there is any way to remove the typing.Any and to statically encode that the returned type must match the type key? Thanks for any advice!

I am using python 3.11.

Solution

Type mappings are not supported in Python (yet?). You can, however, write a Protocol which imitates some aspects of collections.abc.MutableMapping to achieve roughly the same effect, if you're willing to forgo writing it as a direct dictionary (dictionary literal):

import typing as t

T = t.TypeVar("T")

class MyDict(t.Protocol):
    def __getitem__(self, item: type[T]) -> t.Callable[[bytes], T]: ...
    def __setitem__(self, key: type[T], value: t.Callable[[bytes], T]) -> None: ...

# Can't write a dictionary literal with all the entries here, but you can initialise, set, and retrieve values in a type-safe manner
MY_DICT: MyDict = {}
MY_DICT[bytes] = lambda x: x
MY_DICT[int] = lambda x: int.from_bytes(x, "big")
MY_DICT[str] = lambda x: x.decode("utf-8")

# Type-checking in effect by mypy
MY_DICT[bytes] = b""  # mypy: Incompatible types in assignment (expression has type "bytes", target has type "Callable[[bytes], bytes]") [assignment]
reveal_type(MY_DICT[list[int]])  # mypy: Revealed type is "def (builtins.bytes) -> builtins.list[builtins.int]"

If you need more method API from dict, just fill up the method stub signatures in MyDict.