含TypeVar签名的Protocol实现类无法使用自定义类型，如何修正？

我想给Protocol使用者提供自定义任意键值类型数据结构的能力，这类结构通常是Dict，也能适配Redis等外部存储。尝试用Protocol实现时，mypy不允许实现类在方法签名里用自定义类型，相关代码及报错如下：

from typing import (
    Dict,
    Generic,
    Optional,
    Protocol,
    Type,
    TypeVar,
    Union,
    runtime_checkable,
)

Key = TypeVar("Key")
Value = TypeVar("Value")


@runtime_checkable
class Data(Protocol):
    async def get_value(self, key: Key) -> Optional[Value]:
        ...

    async def set_value(self, key: Key, value: Value) -> None:
        ...


class DataUser(Generic[Key, Value]):
    def __init__(
        self,
        cls_data: Type[Data],
    ) -> None:
        self.__data = cls_data()

    async def get_value(self, key: Key) -> Optional[Value]:
        return await self.__data.get_value(key)

    async def set_value(self, key: Key, value: Value) -> None:
        await self.__data.set_value(key=key, value=value)


KVDataValue = Union[str, int, float]


class KVData:
    def __init__(self) -> None:
        self.__data: Dict[str, KVDataValue] = {}

    async def get_value(self, key: str) -> Optional[KVDataValue]:
        return self.__data.get(key)

    async def set_value(self, key: str, value: KVDataValue) -> None:
        self.__data[key] = value


data_user = DataUser[str, KVDataValue](
    cls_data=KVData,
)

运行mypy检查时的错误：

src/aiomon/example.py:57:14: error: Argument "cls_data" to "DataUser" has incompatible type "Type[KVData]"; expected "Type[Data]"  [arg-type]
        cls_data=KVData,
                 ^~~~~~
Found 1 error in 1 file (checked 11 source files)

注：PyCharm能正确识别类型注解，且移除KVData的set_value方法时，IDE能正确识别它不再符合Data Protocol。

问题原因 #

当前的Data Protocol没有和Key、Value两个TypeVar绑定，导致mypy无法将KVData（使用str和KVDataValue作为键值类型）与通用的Data Protocol关联起来。

解决步骤 #

1. 让Data Protocol继承Generic[Key, Value]，明确它是一个泛型协议
2. 修改DataUser的cls_data参数类型，指定为Type[Data[Key, Value]]，让它与当前泛型实例的键值类型匹配

修改后的完整代码 #

from typing import (
    Dict,
    Generic,
    Optional,
    Protocol,
    Type,
    TypeVar,
    Union,
    runtime_checkable,
)

Key = TypeVar("Key")
Value = TypeVar("Value")


@runtime_checkable
class Data(Protocol, Generic[Key, Value]):  # 添加Generic继承，绑定TypeVar
    async def get_value(self, key: Key) -> Optional[Value]:
        ...

    async def set_value(self, key: Key, value: Value) -> None:
        ...


class DataUser(Generic[Key, Value]):
    def __init__(
        self,
        cls_data: Type[Data[Key, Value]],  # 指定泛型Protocol的具体类型
    ) -> None:
        self.__data = cls_data()

    async def get_value(self, key: Key) -> Optional[Value]:
        return await self.__data.get_value(key)

    async def set_value(self, key: Key, value: Value) -> None:
        await self.__data.set_value(key=key, value=value)


KVDataValue = Union[str, int, float]


class KVData:
    def __init__(self) -> None:
        self.__data: Dict[str, KVDataValue] = {}

    async def get_value(self, key: str) -> Optional[KVDataValue]:
        return self.__data.get(key)

    async def set_value(self, key: str, value: KVDataValue) -> None:
        self.__data[key] = value


data_user = DataUser[str, KVDataValue](
    cls_data=KVData,
)

这样修改后，mypy就能正确识别KVData符合Data[str, KVDataValue]的类型约束，不会再抛出类型不兼容的错误。

内容的提问来源于stack exchange，提问作者Lev Rubel