CODEWITHSUNDEEP
pythonpython-3.xinheritancetype-hintingmypy
tingled
tingled

Reputation: 33

Python type annotations: return type of inherited method

I have created a custom dictionary-like class to simplify merging evaluation metrics across a large dataset. This class implements an __add__ method to sum up various metrics.

Here is a simplified version of the code I'm working on:

from __future__ import annotations
from typing import TypeVar, Dict


T = TypeVar('T', int, float)


class AddableDict(Dict[str, T]):
    def __add__(self, other: AddableDict[T]) -> AddableDict[T]:
        if not isinstance(other, self.__class__):
            raise ValueError()
        new_dict = self.__class__()
        all_keys = set(list(self.keys()) + list(other.keys()))
        for key in all_keys:
            new_dict[key] = self.get(key, 0) + other.get(key, 0)
        return new_dict


# AddableIntDict = AddableDict[int]
# this would work just fine, however I need to add a few additional methods


class AddableIntDict(AddableDict[int]):
    def some_int_specific_method(self) -> None:
        pass


def main() -> None:
    x = AddableIntDict()
    y = AddableIntDict()
    x['a'] = 1
    y['a'] = 3

    x += y  # breaks mypy

The final line of the program breaks mypy (0.782) with the following error:

error: Incompatible types in assignment (expression has type "AddableDict[int]", variable has type "AddableIntDict")

This error makes sense to me.

The code works fine when I define AddableIntDict as a type alias of AddableDict[int], as noted in my comment, however because I needed to add additional methods depending on the type of the dictionary's values, as indicated by some_int_specific_method, I can't simply use a type alias.

Could anyone point me in the right direction for how to annotate the parent classes's __add__ method so that it will return the type of the calling class?

(I'm using Python 3.8.3)

Upvotes: 3

Views: 1728

Answers (1)

MisterMiyagi
MisterMiyagi

Reputation: 52089

One can refer to "the type of self" by using a type variable. This resolves to the appropriate type of either the base class or subclass on which the method is invoked:

from typing import TypeVar, Dict


T = TypeVar('T', int, float)
AD = TypeVar('AD', bound='AddableDict')


class AddableDict(Dict[str, T]):
    def __add__(self: AD, other: AD) -> AD: ...


class AddableIntDict(AddableDict[int]):
    def some_int_specific_method(self) -> None: ...

x = AddableIntDict(a=1)
y = AddableIntDict(a=3)
x += y  # works for mypy and others

Upvotes: 6

Related Questions