Return type annotations for command handler that dispatches based on incoming type?

Python Help
1

Hello,

I am creating a class which routes incoming requests (dataclasses) based on their type to a handler which is registered at runtime.

Can someone think of a way to make my IDE (PyCharm) aware of the return type based on the input?
Is this even possible in Python?

This is the minimal sample that I put together. What I’d like to achieve is that my IDE can reveal the type of actual to be str (or whatever is annotated in the handler callable).

import dataclasses
import typing as t


@dataclasses.dataclass
class One:
    ...


@dataclasses.dataclass
class Two:
    ...


class OneHandler:
    def do(self, payload: One) -> str:
        return "ok"


class TwoHandler:
    def do(self, payload: Two) -> int:
        return 5


class Dispatcher:
    def __init__(self) -> None:
        self._handlers: dict[t.Any, t.Callable] = {}

    def register(self, payload_cls: t.Union[t.Type[One], t.Type[Two]],
                 handler: t.Callable):
        self._handlers[payload_cls] = handler

    def handle(self, payload):
        handler = self._handlers[type(payload)]
        return handler(payload)


def test_():
    one_handler = OneHandler()
    two_handler = TwoHandler()
    d = Dispatcher()
    d.register(One, one_handler.do)
    d.register(Two, two_handler.do)

    actual = d.handle(One())
    assert isinstance(actual, str)

Bye,
chbndrhnns

2

Well, you can use @overload for this

    @overload
    def handle(self, payload: One) -> str: ...
    @overload
    def handle(self, payload: Two) -> int: ...
    
    def handle(self, payload: One | Two) -> str | int:
        handler = self._handlers[type(payload)]
        return handler(payload)

Screenshot 2022-12-08 at 13.20.15
Screenshot 2022-12-08 at 13.20.151504×444 51.9 KB

But you have to explicitly encode all the possibilities manually in the overloads, which is pretty unpleasant. Maybe there is a fancier way to automate the type detection with protocols and generics.

3

I was thinking about overload and was wondering if it’s possible to generate the overloads at inspection time for the IDE somehow. I did not want to the route of auto-generating python code just for that purpose.