Get mypy to accept subtype of generic type as a method argument

Question

I'm trying to extract a pattern that we are using in our code base into a more generic, reusable construct. However, I can't seem to get the generic type annotations to work with mypy.

Here's what I got:

from abc import (
    ABC,
    abstractmethod
)
import asyncio
import contextlib
from typing import (
    Any,
    Iterator,
    Generic,
    TypeVar
)

_TMsg = TypeVar('_TMsg')

class MsgQueueExposer(ABC, Generic[_TMsg]):

    @abstractmethod
    def subscribe(self, subscriber: 'MsgQueueSubscriber[_TMsg]') -> None:
        raise NotImplementedError("Must be implemented by subclasses")

    @abstractmethod
    def unsubscribe(self, subscriber: 'MsgQueueSubscriber[_TMsg]') -> None:
        raise NotImplementedError("Must be implemented by subclasses")


class MsgQueueSubscriber(Generic[_TMsg]):

    @contextlib.contextmanager
    def subscribe(
            self,
            msg_queue_exposer: MsgQueueExposer[_TMsg]) -> Iterator[None]:
        msg_queue_exposer.subscribe(self)
        try:
            yield
        finally:
            msg_queue_exposer.unsubscribe(self)


class DemoMsgQueSubscriber(MsgQueueSubscriber[int]):
    pass

class DemoMsgQueueExposer(MsgQueueExposer[int]):

    # The following works for mypy:

    # def subscribe(self, subscriber: MsgQueueSubscriber[int]) -> None:
    #     pass

    # def unsubscribe(self, subscriber: MsgQueueSubscriber[int]) -> None:
    #     pass

    # This doesn't work but I want it to work :)

    def subscribe(self, subscriber: DemoMsgQueSubscriber) -> None:
        pass

    def unsubscribe(self, subscriber: DemoMsgQueSubscriber) -> None:
        pass

I commented out some code that works but doesn't quite fulfill my needs. Basically I want that the DemoMsgQueueExposer accepts a DemoMsgQueSubscriber in its subscribe and unsubscribe methods. The code type checks just fine if I use MsgQueueSubscriber[int] as a type but I want it to accept subtypes of that.

I keep running into the following error.

generic_msg_queue.py:55: error: Argument 1 of "subscribe" incompatible with supertype "MsgQueueExposer"

I feel that this has something to do with co-/contravariants but I tried several things before I gave up and came here.

Answer 1

Your best bets are either to 1) just delete subscribe and unsubscribe from MsgQueueExposer altogether, or 2) make MsgQueueExposer generic with respect to the subscriber, either in addition to or instead of the msg.

Here is an example of what approach 2 might look like, assuming we want to keep the _TMsg type parameter. Note that I added a messages() method for demonstration purposes:

from abc import ABC, abstractmethod
import asyncio
import contextlib
from typing import Any, Iterator, Generic, TypeVar, List

_TMsg = TypeVar('_TMsg')
_TSubscriber = TypeVar('_TSubscriber', bound='MsgQueueSubscriber')

class MsgQueueExposer(ABC, Generic[_TSubscriber, _TMsg]):

    @abstractmethod
    def subscribe(self, subscriber: _TSubscriber) -> None:
        raise NotImplementedError("Must be implemented by subclasses")

    @abstractmethod
    def unsubscribe(self, subscriber: _TSubscriber) -> None:
        raise NotImplementedError("Must be implemented by subclasses")

    @abstractmethod
    def messages(self) -> List[_TMsg]:
        raise NotImplementedError("Must be implemented by subclasses")


class MsgQueueSubscriber(Generic[_TMsg]):
    # Note that we are annotating the 'self' parameter here, so we can
    # capture the subclass's exact type.

    @contextlib.contextmanager
    def subscribe(
            self: _TSubscriber,
            msg_queue_exposer: MsgQueueExposer[_TSubscriber, _TMsg]) -> Iterator[None]:
        msg_queue_exposer.subscribe(self)
        try:
            yield
        finally:
            msg_queue_exposer.unsubscribe(self)


class DemoMsgQueSubscriber(MsgQueueSubscriber[int]):
    pass

class DemoMsgQueueExposer(MsgQueueExposer[DemoMsgQueSubscriber, int]):
    def subscribe(self, subscriber: DemoMsgQueSubscriber) -> None:
        pass

    def unsubscribe(self, subscriber: DemoMsgQueSubscriber) -> None:
        pass

    def messages(self) -> List[int]:
        pass

More broadly, we wanted to express the idea that each MsgQueueExposer works only for a specific kind of subscriber, so we needed to encode that information somewhere.

The one hole in this is that mypy will not be able to make sure when you're using MsgQueueExposer that whatever type the subscriber receives and that whatever type the exposer is expecting will agree. So, if we defined the demo subscriber as class DemoMsgQueSubscriber(MsgQueueSubscriber[str]) but kept DemoMsgQueueExposer the same, mypy would be unable to detect this mistake.

But I'm assuming you're always going to be creating a new subscriber and a new exposer in pairs and is something you can carefully audit, so this mistake is probably unlikely to occur in practice.