[Python] How to use type hints in a dict with TypeVar and a function that depends on that...

I have a relative simple app in which I use a Message Bus to pass messages.

I prefer a bit more stricter type checking: rather than passing arbitrary text strings as topic, I pass message objects, and to be exact: I pass instances of a subclass of the BaseMessage class, and I want to ensure that the handler is designed to take that specific subclass as parameter.

Here is the working code:

from dataclasses import dataclass, field, replace
from typing import Type, TypeVar, Callable

class BaseMessage:
# Abstract class. All messages must be a subclass of BaseMessage
pass

# Message represent any subclass of BaseMessage. Used for type hints
Message = TypeVar('Message', bound=BaseMessage)

# Define the different messages that can be passed around:
@dataclass
class MessageType1(BaseMessage):
data: int

@dataclass
class MessageType2(BaseMessage):
point: tuple[float, float]


class MessageBus():
def __init__(self):
self.subscribers = {}

def subscribe(self, messagetype: Type[Message], handler: Callable[[Message], None]) -> None:
"""Register a handler (subscriber) for a topic."""
self.subscribers.setdefault(messagetype, []).append(handler)

def publish(self, message: BaseMessage):
"""Send a message to all handlers of the topic."""
for handler in self.subscribers.get(type(message), []):
handler(message)


# Example usage
def my_handler1(message: MessageType1):
print(f"my_handler1 message: {message}")

def my_handler2(message: MessageType2):
print(f"my_handler2 message: {message}")

bus = MessageBus()
bus.subscribe(MessageType1, my_handler1)
bus.subscribe(MessageType2, my_handler2)
bus.subscribe(MessageType1, my_handler2) # I want to see an error here somewhere!
bus.publish(MessageType1(3))
bus.publish(MessageType2((5.0, 2.5)))


As you can see, I register a few handlers my_handler1 and my_handler2, and then register them with the message bus, telling upon what message type what function should be called (with that message as parameter).

For illustration purpose, this code contains a bug: bus.subscribe(MessageType1, my_handler2). This is wrong, as my_handler2 takes a MessageType2 instance instead of a MessageType1 instance.

The good news is that in the above code, pyright alerts me about this issue:

error: Argument of type "(message: MessageType2) -> None" cannot be assigned to parameter "handler" of type "(Message@subscribe) -> None" in function "subscribe"
Type "(message: MessageType2) -> None" is not assignable to type "(MessageType1) -> None"
Parameter 1: type "MessageType1" is incompatible with type "MessageType2"
"MessageType1" is not assignable to "MessageType2" (reportArgumentType)


So everything works.

The only thing what I still like to do, is to type-hint MessageBus.subscribers in this snippet:

class MessageBus():
def __init__(self):
self.subscribers = {}


I kind of assumed that this would be either

self.subscribers: dict[Type[Message], list[Callable[[Message], None]]] = {}


or

self.subscribers: dict[Type[BaseMessage], list[Callable[[BaseMessage], None]]] = {}


However, with the first, pyright gives an error Type Variable "Message" has no meaning in this context. And with the second line, pyright gives an error during during the append() call: Type "(Message@subscribe) →> None" is not assignable to type "(BaseMessage) →> None"

So my questions:

1. What is the correct type hint for self.subscribers in the above code? (and why?)
2. Would it be possible to define a type hinting shortcut for the function call, to enhance readability? E.g. MessageHandler = Callable[[Message], None]
3. Alternatively, is there a recommended solution with a subscriber-publish model that passes Message Objects? I prefer solutions that use the Python standard library instead of external packages.

Continue reading...