How to properly type Python methods returning a dict with a generically typed instance and its generically typed property?

user_5921935 0
Viewed 272

I have the following Python function and want to type it in a way that mypy understands the return type as a dictionary, with the items' type as the value type and the id type as the key type:

def id_map(items):
    return {item.id: item for item in items}

The current approach looks like this:

from typing import Iterable, Protocol, TypeVar

S = TypeVar("S", covariant=True, bound="ClassWithId")
T = TypeVar("T", covariant=True)

class ClassWithId(Protocol[S, T]):
    def __new__(cls: type[S], *args: Any, **kwargs: Any) -> "ClassWithId[S, T]":
        ...

    @property
    def id(self: S) -> T:
        ...


def id_map(items: Iterable[ClassWithId[S, T]]) -> dict[T, S]:
    return {item.id: cast(S, item) for item in items}

This currently results in the type builtins.dict[Any, Any]. I cannot think of a different implementation that would make that work. I am wondering whether this is possible at all using mypy.

python-typing mypy python-3.x python
1 Answers

TLDR: Remove the self type variable S. It is not needed but leads to a type relation that is unlikely to be understood by type checkers.

class ClassWithId(Protocol[T]):
    @property
    def id(self) -> T:
        ...

Type variables can be used in two flavours: as inferred variables in a call or as declared variables in a generic type. Use the former for types that are constrained from usage context, and the latter for types that must be consistent across an entire class.

The S type var is a Self type; this can always be inferred from usage and does not need to be repeated in the type itself. Otherwise, you construct a recursive type which is likely to confuse current type checkers.

Removing S from the generic/protocol shows that it is not actually used for any type relation: The T is fully constrained by the class ClassWithId[T]. There is no need to refer to S to define any relation between input and output.

class ClassWithId(Protocol[T]):
    @property
    def id(self: S) -> T: ...
    #                  ^ defined by generic class, not S

Notice how id_map also does not depend on S – knowing either of T or ClassWithId[T] is enough to know the other. As a result, one can remove S completely.

T = TypeVar("T", covariant=True)

class ClassWithId(Protocol[T]):
    @property
    def id(self) -> T:
        ...

def id_map(items: Iterable[ClassWithId[T]]) -> dict[T, ClassWithId[T]]:
    return {item.id: item for item in items}

If id_map is supposed to preserve the exact ClassWithId[T] and not just the protocol, this corresponds to a higher kinded type. In principle, this is expressed by a Type variable with a bound of the protocol.

CWI = TypeVar("CWI", bound=ClassWithId)

def id_map(items: Iterable[CWI[T]]) -> dict[T, CWI[T]]:
    return {item.id: item for item in items}

However, MyPy does not supported higher kinded types at the moment.

user_53499160
Related