Type hints support for subclass of dict

Question

How can I implement a subclass of dict so it supports type hints like vanilla dict?

I want to create a custom dict that does some extra work when modifying item, I implement it by subclassing the built-in dict class. After doing some googling I learned that I need to use the generic class from typing to achieve this:

_KT = TypeVar("_KT")
_VT = TypeVar("_VT")

class CustomDict(dict, Mapping[_KT, _VT]):
    """A dict that call callback function after setting or popping item."""

    def __init__(self):
        self.callback = None

    def __setitem__(self, key, value):
        super().__setitem__(key, value)
        if self.callback:
            self.callback()

    def pop(self, key, *arg):
        super().pop(key, *arg)
        if self.callback:
            self.callback()

    def register_callback(self, callback: Callable[[], None]):
        self.callback = callback

For a vanilla dict object, type annotation is added like:

my_dict: dict[int, str] = {}

Now my custom dict supports type annotation just like the vanilla dict:

my_custom_dict: CustomDict[int, str] = CustomDict()

But mypy doesn’t throw me an error if I add item with incompatible type:

# the key is type str and the value is type int, which is against the annotation
my_custom_dict["3.234"] = 2

For a vanilla dict, doing this mypy will give me an error. So which part did I get wrong and how can I implement the subclass along with the annotation correctly? Any help would be appreciated!

Asked By: oeter

||

Source

Answer 1

You didn’t actually use your type variables in the definition of __setitem__, which means it isn’t typechecked. (Your also forgot to return the value of super().pop(key, *args).) Add the type hints:

from typing import TypeVar, Mapping, Callable
_KT = TypeVar("_KT")
_VT = TypeVar("_VT")

class CustomDict(dict, Mapping[_KT, _VT]):
    """A dict that call callback function after setting or popping item."""

    callback: Optional[Callable[[], None]]

    def __init__(self):
        self.callback = None

    def __setitem__(self, key: _KT, value: _VT):
        super().__setitem__(key, value)
        if self.callback:
            self.callback()

    def pop(self, key, *arg) -> _VT:
        rv = super().pop(key, *arg)
        if self.callback:
            self.callback()
        return rv

    def register_callback(self, callback: Callable[[], None]):
        self.callback = callback

my_custom_dict: CustomDict[int, str] = CustomDict()

my_custom_dict["3.234"] = 2

and you’ll get the expected errors:

tmp.py:27: error: Invalid index type "str" for "CustomDict[int, str]"; expected type "int"  [index]
tmp.py:27: error: Incompatible types in assignment (expression has type "int", target has type "str")  [assignment]
Found 2 errors in 1 file (checked 1 source file)

Note that you don’t need to inherit from Mapping (or MutableMapping, which is what you would inherit from if you want to support __setitem__, though inheriting from dict as well is sufficient). Just use dict generically:

from typing import TypeVar, Callable
_KT = TypeVar("_KT")
_VT = TypeVar("_VT")

class CustomDict(dict[_KT, _VT]):
    ...

As pointed out in a comment, consult typeshed (the repository of Python stub files for built-in and standard-library code) to see the types of the methods you are overriding, either so you can copy them or ensure that your custom type hints are consistent with them.

Unrelated, consider using the do-nothing function lambda: None as the default callback instead of None. Then you can simply call the callback without checking if it exists.

Answered By: chepner

Type hints support for subclass of dict

Question:

Answers: