CODEWITHSUNDEEP
python
Shai Avr
Shai Avr

Reputation: 1350

How to add type annotations to custom dict subclass in python?

I have a custom dict subclass that is similar to defaultdict but passes the missing key to default_factory so it can generate an appropriate value.

class KeyDefaultDict(dict):
    __slots__ = ("default_factory",)

    def __init__(self, default_factory, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.default_factory = default_factory

    def __missing__(self, key):
        if self.default_factory is None:
            raise KeyError(key)
        ret = self[key] = self.default_factory(key)
        return ret

    def __repr__(self):
        return (
            f"{type(self).__name__}({repr(self.default_factory)}, {super().__repr__()})"
        )


d = KeyDefaultDict(int)
print(d["1"] + d["2"] + d["3"])  # out: 6
print(d)  # out: KeyDefaultDict(<class 'int'>, {'1': 1, '2': 2, '3': 3})

I wanted to add type annotations for this class like the rest of my project, but I couldn't find any example of how to do that. I saw that the typing module uses external classes for adding annotations. For example, defaultdict would be annotated with typing.DefaultDict whose definition is class typing.DefaultDict(collections.defaultdict, MutableMapping[KT, VT]).

So it's an external class that subclasses defaultdict and the generic typing.MutableMapping. However, I thought that they probably did it this way because they didn't want to change the original collections.defaultdict. I found examples of subclasses of Generic and Mapping but not classes that inherit from something else like dict.

The question is: How can I add type annotations for this class to make it a generic class? Do I need to extend something else or make an external class for annotations?

I am using python 3.7.5 and I prefer to inherit directly from dict so I don't have to implement the required methods and for performance reasons.

Thanks in advance.

Upvotes: 5

Views: 2675

Answers (1)

Alex
Alex

Reputation: 177

I am very late but I just got this working in my own code base.

Basically, you need to use the typing Mapping generic This is the generic of what dict uses so you can define other types like MyDict[str, int].

For my use-case, I wanted a special dict that formatted itself cleanly for logging but I will be using it all over with various types so it needed typing support.

How to:

import typing

# these are generic type vars to tell mapping to accept any type vars when creating a type
_KT = typing.TypeVar("_KT") #  key type
_VT = typing.TypeVar("_VT") #  value type


# `typing.Mapping` requires you to implement certain functions like __getitem__
# I didn't want to do that, so I just subclassed dict.
# Note: The type you're subclassing needs to come BEFORE 
# the `typing` subclass or the dict won't work.
# I had a test fail where my debugger showed that the dict had the items,
# but it wouldn't actually allow access to them

class MyDict(dict, typing.Mapping[_KT, _VT]):
    """My custom dict that logs a special way"""
    def __str__(self):
        # This function isn't necessary for your use-case, just including as example code
        return clean_string_helper_func(
            super(MyDict, self).__str__()
        )

# Now define the key, value typings of your subclassed dict
RequestDict = MyDict[str, typing.Tuple[str, str]]
ModelDict = MyDict[str, typing.Any]

Now use you custom types of your subclassed dict:

from my_project.custom_typing import RequestDict #  Import your custom type

request = RequestDict()
request["test"] = ("sierra", "117")

print(request)

Will output as { "test": ("sierra", "117") }

Upvotes: 4

Related Questions