使用 Protocol 和 TypeVar 来指定任意数据类的 Python 类型提示

Question

我正在编写一个可以在内存中存储任意数据类的类。我试图指定要存储的实例必须是数据类并且具有

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

有人可以帮助我了解类型提示吗？

最佳 拉尔斯

Answer 1

MisterMiyagi 向我指出了 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

 一起使用的 

typing.Protocol

字段