Implementers of a Protocol with TypeVars in signatures can't use their own types instead

Question

I'm trying to give users of my protocols a way to use whatever they want as types for keys and values of some arbitrary data structure.

Usually it will be a Dict, but it can easily be some external data storage like Redis.

I tried using Protocol to give this options to users, but mypy doesn't allow implementers to use their custom defined types in methods singnatures.

Is there a way to fix this typing or am I fundamentally using wrong tools here?

I'm using Python 3.11 and here's the script:

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,
)

Running mypy on this code results in:

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)

Worth noting that I'm using PyCharm as my IDE and it shows that this is correct typing. Moreover when I, for example, remove set_value method from KVData it correctly identifies that this class stopped adhering to the Data Protocol.

Answer 1

If the protocol is supposed to be generic in terms of what arguments its methods accept (and return), you should define it as such. (see "Generic Protocols" in PEP 544)

After all, DataUser is already generic and (if I understand your use case correctly) its type parameters should be bound to those of the underlying Data class.

This seems to work just fine:

from typing import Generic, Optional, Protocol, TypeVar, Union, runtime_checkable

K = TypeVar("K", contravariant=True)
V = TypeVar("V")


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

    async def set_value(self, key: K, value: V) -> None: ...


class DataUser(Generic[K, V]):
    def __init__(self, cls_data: type[Data[K, V]]) -> None:
        self.__data = cls_data()

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

    async def set_value(self, key: K, value: V) -> 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)

Passes mypy --strict without errors.