Search code examples
pythondictionaryoverloadingtype-hintingpython-typing

Type hint for special keys of dict subclass


I would like to create type hints for a specific key of a dictionary. That is to say that the value of a key with a specific name has a specific type.

Say we have a dict with value type Union[int,str] and now want to specify that only the key "intkey" has a corresponding value of type int and all other keys have a value type str.

How can we do this? My attempt is below, but that does not work with mypy.

Note: We cannot use TypedDict here because with that all keys must be known in advance.

class SpecialDict(Dict[str,Union[int,str]]):
    @overload                                      # <---- error A
    def __getitem__(self, key: Literal["intkey"],/) -> int: ...
    @overload
    def __getitem__(self, key: str,/) -> str: ...

    def __getitem__(self, key,/): return super().__getitem__(key)

    @overload                                       # <---- error B
    def __setitem__(self, key: Literal["intkey"], val: int,/) -> None: ...
    @overload
    def __setitem__(self, key: str, val: str, /) -> None: ...

    def __setitem__(self, key, val,/): return super().__setitem__(key, val) 

But that does not work. mypy 1.4 generates the following errors

  • A
    Overloaded function signatures 1 and 2 overlap with incompatible return types  [misc]
    
  • B
    Signature of "__setitem__" incompatible with supertype "dict"  [override] 
       Superclass:
           def __setitem__(self, str, Union[int, str], /) -> None 
       Subclass:
           @overload
           def __setitem__(self, Literal['intkey'], int, /) -> None
           @overload
           def __setitem__(self, str, str, /) -> None 
    Signature of "__setitem__" incompatible with supertype "MutableMapping"  [override] 
       Superclassmypy(note)
           def __setitem__(self, str, Union[int, str], /) -> None 
       Subclass:
           @overload
           def __setitem__(self, Literal['intkey'], int, /) -> None 
           @overload
           def __setitem__(self, str, str, /) -> None
    

I don't know why these errors are there and not even the slightest idea what the problem is of error B.


Solution

  • Error A: Unsafely overlapping overloads

    Thankfully, this exact type of error is very well explained in the Mypy docs. Allow me to adapt their example to more closely fit your situation, but simplify it a bit and lose the surrounding class:

    from typing import Literal, Union, overload
    
    
    @overload
    def f(a: Literal[""]) -> int: ...
    @overload
    def f(a: str) -> str: ...
    
    
    def f(a: str) -> Union[int, str]:
        if a == "":
            return 42
        return a.upper()
    

    The issue here is that Literal[""] is a subtype of str and your signatures imply that when something of that type (i.e. the literal empty string "") is passed to f, it will return an integer, but when the supertype is passed it will return a string.

    Now imagine the following:

    s: str
    ...
    s = ""
    print(f(s) + " some string")  # TypeError!
    

    In other words, s can hold any string and at runtime at some point it is assigned the empty string. Now, once f called with s as the argument, f will see that it is "" and return an integer, but the type checker can only see that s is declared to be str and thus assume the return type to be str as well. Therefore it has no way of warning you about the coming error attempting to add a string to an integer.

    As the documentation puts it, if it were to allow such an overloaded signature, it would

    result in a discrepancy between the inferred type and the actual runtime type when we try using it [...]

    Note that this should be considered a design decision of the people writing the type checker. They decide how strictly they want to apply their rules for safe/unsafe overloads. Mypy does not apply this rule consistently/to its last logical conclusion and they say this clearly in the documentation. But they also say this: (emphasis mine)

    What this does mean [..] is that you should exercise caution when designing or using an overloaded function that can potentially receive values that are an instance of two seemingly unrelated types.

    In other words, if you are aware of these dangers, I would say it is perfectly reasonable to apply a type: ignore[mics] directive next to your first overloaded signature and move ahead.


    Error B: Violating the Liskov Substitution Principle

    This is more easily explained, but you have to really familiarize yourself with the . The implications are also documented by Mypy.

    You are defining a subtype of dict and overriding one of its methods. The rules of the LSP dictate:

    It’s unsafe to override a method with a more specific argument type [or] with a more general return type.

    Phrased differently in terms of variance, subtyping is contravariant in the method parameter types.

    Since you passed str as the key type argument and Union[int, str] as the value type argument for dict during inheritance, this collapses the parameter types in all its methods str and Union[int, str] for the keys and values respectively. This includes the __setitem__ method.

    But in none of your overload signatures, can you call __setitem__ with those types. You are strengthening the preconditions for that method by only allowing either

    1. a subtype of str (namely Literal["intkey"]) for the key and a subtype of Union[int, str] (namely int) for the value or
    2. str for the key (which is fine) and a subtype of Union[int, str] (namely str) for the value.

    You explicitly do not allow any str for the key and str | int for the value in that method. And strengthening preconditions violates the LSP, ergo is not safe subtyping.

    As a practical example, imagine the following:

    def func(dictionary: dict[str, Union[int, str]]) -> None:
        dictionary["intkey"] = "spam"
    
    
    d = SpecialDict()
    d["intkey"] = 2
    
    func(d)
    print(d["intkey"] ** 2)  # TypeError!
    

    This is what what would pass a static type check without a problem, if SpecialDict were considered a safe subtype of dict[str, Union[int, str]].

    There is no real way around this issue in my opinion, if you really want to subclass dict like this, without resorting to a deliberate type: ignore[override] directive and "hoping it all works out" at runtime.

    An alternative would be to subclass the more general Mapping ABC, because that does not have a __setitem__ method at all and just add that (and any other missing mix-in methods you need like pop) yourself. Then there will obviously be no override issue. But that would also mean your SpecialDict will not be considered a dict or a even a MutableMapping subclass by type checkers.