How to type a generic callable in python

Question

I am trying to do something like this:

from collections.abc import Callable, Coroutine
from typing import Any, Generic, TypeVar

CRT = TypeVar("CRT", bound=Any)

class Command(Generic[CRT]): pass

CommandHandler = Callable[[Command[CRT]], Coroutine[Any, Any, CRT]]

class MyCommand(Command[int]): pass

async def my_command_handler(command: MyCommand) -> int:
    return 42

async def process_command(command: Command[CRT], handler: CommandHandler[CRT]) -> CRT:
    return await handler(command)

async def main() -> None:
    my_command = MyCommand()

    await process_command(my_command, my_command_handler)

But I get the following mypy error on my await process_command(my_command, my_command_handler) line

Argument 2 to "process_command" has incompatible type "Callable[[MyCommand], Coroutine[Any, Any, int]]"; expected "Callable[[Command[int]], Coroutine[Any, Any, int]]"Mypyarg-type
(function) def my_command_handler(command: MyCommand) -> CoroutineType[Any, Any, int]

Which I don't understand because MyCommand is a Command[int].

How do I make mypy happy here?

My end goal is to be able to have different command classes, each command class defines the expected result of the command, and then type some function that takes both the command class, and a handler for the command class, and ensure that the handler takes an instance of the command as a parameter and returns a the correct result defined by the command.

If I change my_command_handler this:

async def my_command_handler(command: Command[int]) -> int:
    if not isinstance(command, MyCommand):
        raise TypeError("Expected MyCommand")
    return 42

It type checks fine, but I'm not a fan of having to do the extra type check in the function to get proper typing inside the function.

Answer 1

You want to have command be compatible with the handler's input. Because of the discussed contravariance is MyCommand and Command[int] not compatible here.

A solution is to disentangle this and use two type variables to link both together:

CommandT = TypeVar("CommandT", bound="Command")

async def process_command2(command: CommandT, handler: HandlerType[CommandT, CRT]) -> CRT:
    return await handler(command)

However as this does not bind CRT to the command and a handler could also convert the type, i.e. a handler of this type would also be valid:

async def unexpected_handler(o: Command[int]) -> str:
   ...

At least for mypy, I found that defining your handler this way:

HandlerType = TypeAliasType("HandlerType", 
    Callable[[CommandT], Coroutine[Any, Any, CRT]]
   | Callable[[Command[CRT]], Coroutine[Any, Any, CRT]],
type_params=(CommandT, CRT,))

does indeed make unexpected_handler incompatible.

async def main() -> None:
    my_command = MyCommand()

    await process_command(my_command, my_command_handler)  # error
    await process_command2(my_command, my_command_handler) # OK
    await process_command2(my_command, unexpected_handler) # error


# incompatible type 
Callable[[Command[int]], Coroutine[Any, Any, str]];
expected 
Callable[[MyCommand], Coroutine[Any, Any, Never]] | Callable[[Command[Never]], Coroutine[Any, Any, Never]]

Mypy Playground Example

Note: pyright accepts it. From what I know does mypy handle inference a bit differently for function arguments, but so far I cannot explain the exact logic here.

Answer 2

The way I ended up solving this is to create a wrapper type, which I called CommandSpec for the command and return type. This all seems to type check properly.

from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from typing import Any, TypeVar, cast

_Command = TypeVar("_Command", bound=Any)
_Return = TypeVar("_Return", bound=Any)
_CommandHandler = Callable[[_Command], Coroutine[Any, Any, _Return]]


@dataclass(frozen=True, kw_only=True)
class CommandSpec[_Command, _Return]:
    command: _Command


handlers: dict[type[CommandSpec[Any, Any]], _CommandHandler[Any, Any]] = {}


def register_handler(
    command_spec: type[CommandSpec[_Command, _Return]], handler: _CommandHandler[_Command, _Return]
) -> None:
    handlers[command_spec] = handler


async def execute_command(command_spec: CommandSpec[_Command, _Return]) -> _Return:
    command_spec_type = type(command_spec)
    handler = handlers.get(command_spec_type)
    if handler is None:
        raise ValueError(f"No handler registered for {command_spec_type}")

    rv = await handler(command_spec.command)
    return cast(_Return, rv)


@dataclass(frozen=True, kw_only=True)
class MyCommandSpec(CommandSpec["MyCommandSpec.Command", int]):
    class Command: pass


async def my_command_handler(command: MyCommandSpec.Command) -> int:
    return 42


async def main() -> None:
    register_handler(MyCommandSpec, my_command_handler)

    result_int = await execute_command(
        MyCommandSpec(
            command=MyCommandSpec.Command(),
        )
    )
    print(f"Result (int): {result_int}")


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())