Python type hints with Protocol and TypeVar to specify arbitrary dataclasses

Question

I'm writing a class that can store arbitrary dataclasses in memory. There I'm trying to specify that instances to be stored must be a dataclass and have a id field. Also it should be possible to get instances by specifying the classes and the id of the instance.

I'm struggling to define proper type hints. I already figured out, I (probably) need a combination of TypeVar and Protocol. This is my current code:

import typing
import uuid
from collections import defaultdict
from dataclasses import field, dataclass


class DataclassWithId(typing.Protocol):
    __dataclass_fields__: typing.Dict
    id: str


Klass = typing.TypeVar("Klass", bound=DataclassWithId)


class InMemoryDataClassStore:
    def __init__(self):
        self._data_store = defaultdict(lambda: dict())

    def add(self, instance: Klass):
        store_for_class = self._get_store_for_class(instance.__class__)
        store_for_class[instance.id] = instance

    def get(self, klass: typing.Type[Klass], id_: str) -> Klass:
        return self._get_store_for_class(klass)[id_]

    def get_all(self, klass) -> typing.List[Klass]:
        return list(self._get_store_for_class(klass).values())

    def _get_store_for_class(
        self, klass: typing.Type[Klass]
    ) -> typing.Dict[str, Klass]:
        return self._data_store[klass]


auto_uuid_field = field(default_factory=lambda: str(uuid.uuid4()))

@dataclass
class ClassA:
    name: str
    id: str = auto_uuid_field


store = InMemoryDataClassStore()
instance_a = ClassA(name="foo")
store.add(instance_a)
print(store.get(klass=ClassA, id_=instance_a.id).name)
print(store.get(klass=ClassA, id_=instance_a.id).other_name)  # supposed to cause a typing error

if I run mypy against this file, I get

in_memory_data_store.py:45: error: Value of type variable "Klass" of "add" of "InMemoryDataClassStore" cannot be "ClassA"
in_memory_data_store.py:46: error: Value of type variable "Klass" of "get" of "InMemoryDataClassStore" cannot be "ClassA"
in_memory_data_store.py:47: error: Value of type variable "Klass" of "get" of "InMemoryDataClassStore" cannot be "ClassA"
in_memory_data_store.py:47: error: "ClassA" has no attribute "other_name"  # expected

Could please someone help me out about the type hints?

Best Lars

Answer 1

MisterMiyagi pointed me to the mypy issue tracker on Github, where it states that Protocol cannot match dataclasses: https://github.com/python/mypy/issues/6568

class WithId(typing.Protocol):
    id: str


Klass = typing.TypeVar("Klass", bound=WithId)

By simply removing the __dataclass_fields__ from the typing.Protocol subclass, everything works as expected. Actually for my code it doesn't matter whether it's a dataclass. It just needs an id field which works with typing.Protocol