Use MyPy with typing.DefaultDict with fixed factory function

Viewed 1606

I want to do something like this:

from functools import partial
from typing import DefaultDict

def five() -> int:
    return 5  # whatever.

MyDefaultDict = partial(DefaultDict[str, int], five)

my_mapping: MyDefaultDict = MyDefaultDict(x=4)

print(my_mapping['x'])
print(my_mapping['y'])
  • The new type/class thing MyDefaultDict is declared once and has a single name.
  • MyDefaultDict can be used like a class constructor, and its signature is the same as dict's, the only difference being that what you get is a defaultdict with five as its factory.
  • MyPy should understand MyDefaultDict to be the same as (or, preferably, a sub-class of) DefaultDict[str, int]. This is the part that doesn't work in the above code.
  • Unnecessary for my current use, but a nice-to-have in the general case: MyDefaultDict should actually have class properties (e.g. static methods) inherited from defaultdict.

It's not surprising that partial isn't suitable. partial plays poorly with MyPy in general, and at best it would be understood as returning some kind of Callable, which wouldn't be a sensible type for my_mapping. (The actual failure reported is "scratch.py:9: error: Variable "scratch.MyDefaultDict" is not valid as a type".)

A variety of other things seem like they could work, but I haven't been able to get them to:

Subclass and override __init__

class MyDefaultDict(DefaultDict[str, int]):
    def __init__(self, *args, **kwargs):
        DefaultDict[str, int].__init__(self, five, *args, **kwargs)

Passes MyPy.

File "scratch.py", line 9, in __init__
DefaultDict[str, int].__init__(self, five, *args, **kwargs)
TypeError: __init__() got an unexpected keyword argument 'x'

Subclass and override __new__

class MyDefaultDict(DefaultDict[str, int]):
    def __new__(cls, *args, **kwargs):
        return super().__new__(cls, five, *args, **kwargs)

Passes MyPy.
"KeyError: 'y'"; the factory function isn't being used.

Wrap partial in something to make it look like a type

I tried type with three args, typing.NewType, and types.new_class. As far as I can tell this isn't a problem any of them are intended to solve.

I looked into metaclasses briefly and didn't see how they might help with this problem.

I'm using python 3.8, but answers for other versions might be workable.

Does anyone know the right way to do this?

1 Answers

While writing the above I realized that different ways of overriding __init__ give different results.

This also doesn't work:

class MyDefaultDict(DefaultDict[str, int]):
    def __init__(self, *args, **kwargs):
        super().__init__(self, five, *args, **kwargs)

File "scratch.py", line 9, in __init__
super().__init__(self, five, *args, **kwargs)
TypeError: first argument must be callable or None

This does work; it would be nice to understand why:

class MyDefaultDict(DefaultDict[str, int]):
    def __init__(self, *args, **kwargs):
        super().__init__(five, *args, **kwargs)
Related