CODEWITHSUNDEEP
pythonmypypython-typing
Thomas Mailund
Thomas Mailund

Reputation: 1826

Type annotation for generic (protocol) class argument

I'm trying to work out how to add a type annotation for a function argument that should be a class implementing a generic protocol.

As an example, assume I have a protocol for a set that could look something like this:

from typing import (
    Protocol, TypeVar, Iterable
)

T = TypeVar('T', contravariant=True)


class Set(Protocol[T]):
    """A set of elements of type T."""

    def __init__(self, init: Iterable[T]) -> None:
        """Initialise set with init."""
        ...

    def __contains__(self, x: T) -> bool:
        """Test if x is in set."""
        ...

    def add(self, x: T) -> None:
        """Add x to the set."""
        ...

    def remove(self, x: T) -> None:
        """Remove x from the set."""
        ...

and I have an algorithm that uses sets of various types, that I want to parameterise with the set implementation. For simplicity I'll just create a list in this function to use as an example:

from typing import Type

def foo(set_type: Type[Set]) -> None:
    """Do clever stuff."""
    x = list(range(10))
    s = set_type(x)
    ...

Here, mypy tells me that Set is missing a type parameter, which I suppose is correct, but I don't want to give it one, as I plan to use set_type with different types.

If I give Set a TypeVar instead

def foo(set_type: Type[Set[T]]) -> None:
    """Do clever stuff"""
    x = list(range(10))
    s = set_type(x)
    ...

I instead get the warning that I set_type() gets an incompatible type, List[int] instead of Iterable[T], which again is correct, but doesn't help me much.

Is there a way to specify that my function argument can be used as a generic constructor for sets of different types?

Upvotes: 1

Views: 2279

Answers (1)

joel
joel

Reputation: 7877

Protocol says nothing about the signature of __init__, even if it's defined on the Protocol. Type does a similar thing - even if Set isn't a Protocol, Type[Set] says nothing about how the type is called.

I initially suggested using Callable[[Iterable[T]], Set[T]]. However, this is problematic, and only works because I omitted the generic parameter, essentially making it Any, as discussed in this Github issue. You can instead use a (rather verbose) protocol.

class MkSet(Protocol):
    def __call__(self, it: Iterable[T]) -> Set[T]:
        ...

def foo(set_type: MkSet) -> None:
    ...

Upvotes: 3

Related Questions