Search code examples
pythonpython-typingtypeddict

How can I type hint a dictionary where the key is a specific tuple and the value is known?


How can I type hint a dictionary where the key is a specific tuple and the value is known?

For example I want to type hint a dict like this:

class A:
    pass

class B:
    pass

class_map: = {
    (str,): A
    (int,): B
}

some_cls = class_map[(str,)]

The use case will be to go from a known set of bases to a class that was previously defined using those bases.


Solution

  • One can do this by

    • making a new class ClassMap which will allow dict key lookup
      • note: dict cannot be subclassed because our __getitem__ signature differs from the one that dict uses
    • in ClassMap define __getitem__ which gets values from an input dict
    • in ClassMap define an overload of __getitem__ with the tuple input and type hint output
    • making an instance of the ClassMap
    • use it and type hints work
    • passes mypy check

    Another thing that we can do is require that the input to make the dict is a frozenset of tuples. Then one can type hint what is allowed in:

    tuple_items: frozenset[
        typing.Union[
            typing.Tuple[typing_extensions.Literal['1invalid'], int],
            typing.Tuple[typing_extensions.LiteralString, float]
        ]
    ] = frozenset({
        ('1invalid', 1),
        ('a', 1.234)
    })
    

    This enables functionality similar to TypedDict but with tuple keys.

    import typing
    
    class A:
        pass
    
    
    class B:
        pass
    
    
    class ClassMap:
        def __init__(self, data: dict):
            self.__data = data
    
        @typing.overload
        def __getitem__(self, name: typing.Tuple[typing.Type[str]]) -> typing.Type[A]: ...
    
        @typing.overload
        def __getitem__(self, name: typing.Tuple[typing.Type[int]]) -> typing.Type[B]: ...
    
        def __getitem__(self, name):
            return self.__data[name]
    
    class_map = ClassMap({
        (str,): A
        (int,): B
    })
    
    
    some_cls = class_map[(str,)]  # sees A, works!