带有Protocol和TypeVar的Python类型提示可指定任意数据类

问题描述

我正在编写一个可以在内存中存储任意数据类的类。我试图指定要存储的实例必须是数据类并且具有id字段。此外，应该可以通过指定实例的类和ID来获取实例。

我正在努力定义正确的类型提示。我已经弄清楚，我（可能）需要TypeVar和Protocol的组合。这是我当前的代码：

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

如果我对此文件运行mypy，则会得到

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

请有人帮我解决有关类型提示的问题吗？

最佳 拉斯

解决方法

宫城先生将我指向Github上的mypy问题跟踪器，其中指出Protocol无法匹配数据类：https://github.com/python/mypy/issues/6568

class WithId(typing.Protocol):
    id: str


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

只需从__dataclass_fields__子类中删除typing.Protocol，一切都会按预期进行。实际上，对于我的代码，它是否是数据类都没有关系。它只需要一个id的{​​{1}}字段