Python: How to specify function return type in a mixin method

Question

I'm trying to accurately annotate the type of return from a mixin method.

Minimal example below:

from dataclasses import dataclass, field
from typing import Dict, Protocol

class HasStore(Protocol):
    store: dict


class StoreGetMixin:
    def __getitem__(self: HasStore, key: str):  # what to put in -> type?
        return self.store[key]


@dataclass
class ThingOne:
    size: int


@dataclass
class ThingTwo:
    name: str


@dataclass
class Ones(StoreGetMixin):
    store: Dict[str, ThingOne] = field(default_factory=dict)


@dataclass
class Twos(StoreGetMixin):
    store: Dict[str, ThingTwo] = field(default_factory=dict)


ones = Ones()
ones.store = {"a": ThingOne(1), "b": ThingOne(2)}
one = ones["a"]  # <- this should be typed as a ThingOne

twos = Twos()
twos.store = {"a": ThingTwo("one"), "b": ThingTwo("two")}
two = twos["a"]  # <- this should be typed as a ThingTwo

So I'm trying to write a reusable mixin that can access the store dictionary in various dataclasses. But I don't know how to tell the __getitem__ method what type it is returning.

I've seen examples of typing self but i need to type self.store in this case. mypy can see the type of self.store as dict (I assume this comes from HasStore) but it doesn't know the key or value types.

Is this possible?

Do I need to annotate the HasStore protocol's store better? I tried, eg, dict[str, T] but I couldn't make it work because T isn't instantiated if used like that.

The reason I'm doing this is to use Mashumaro to serialise my dataclasses into JSON and I couldn't figure out how to store a plain dictionary keyed off a str without making a single attribute in a dataclass called store. So there are a few of these and they are all have a dict called store so I thought I'd use collections.abc.MutableMapping to allow me to access store via a subscript, eg, twos["a"]. Therefore, I wrote the mixin to do this but it's a pity to lose all the types if I don't need to.

Thanks heaps for any pointers :)

user_4111130 · Answer

A Generic Protocol should work (doc here: https://mypy.readthedocs.io/en/stable/generics.html#generic-protocols):

T = TypeVar('T')


class HasStore(Protocol[T]):
    store: Dict[str, T]


class StoreGetMixin(HasStore[T]):

    def __getitem__(self, key: str) -> T:  # what to put in -> type?
        return self.store[key]


@dataclass
class ThingOne:
    size: int


@dataclass
class ThingTwo:
    name: str


@dataclass
class Ones(StoreGetMixin[ThingOne]):
    store: Dict[str, ThingOne] = field(default_factory=dict)


@dataclass
class Twos(StoreGetMixin[ThingTwo]):
    store: Dict[str, ThingTwo] = field(default_factory=dict)


ones = Ones()
ones.store = {"a": ThingOne(1), "b": ThingOne(2)}
one = ones["a"]  # <- this should be typed as a ThingOne
reveal_type(one)
twos = Twos()
twos.store = {"a": ThingTwo("one"), "b": ThingTwo("two")}
two = twos["a"]  # <- this should be typed as a ThingTwo
reveal_type(two)

output is:

note: Revealed type is 'experiment.ThingOne*'
note: Revealed type is 'experiment.ThingTwo*'

Python: How to specify function return type in a mixin method

1 Answers