9

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

lmr2391
  • 571
  • 1
  • 7
  • 14
  • 1
    Does this answer your question? [type hint for an instance of a non specific dataclass](https://stackoverflow.com/questions/54668000/type-hint-for-an-instance-of-a-non-specific-dataclass) – MisterMiyagi Sep 05 '20 at 12:15
  • As the linked Q&A notes, [``Protocol`` currently cannot match dataclasses](https://github.com/python/mypy/issues/6568). – MisterMiyagi Sep 05 '20 at 12:16
  • Thanks MisterMiyagi for pointing me to the issue on the mypy tracker. For what I'm trying to do, it doesn't really matter that the classes stored are `dataclasses`. So by just removing the `__dataclass_fields__` everything works as expected – lmr2391 Sep 05 '20 at 13:10
  • I think this is a particularly well-constructed question, thanks for posting. – Jason R Stevens CFA Sep 05 '20 at 13:18

1 Answers1

4

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

lmr2391
  • 571
  • 1
  • 7
  • 14
  • 1
    Has anyone gotten this to work with frozen dataclasses? I tried this implementation, but once I make the dataclass frozen, mypy says the type I'm supplying is inadequate. "Value of type "ProtocolType" of foo cannot be FrozenDataclass". I defined a function called "foo" and tried to pass an instance of the "FrozenDataclass" to it. "ProtocolType" is the type that bounded by the protocol. – jaywhy13 Jan 05 '22 at 16:24