Annotate an instance method with metadata without breaking type safety

[GDYendell]

I am creating a control system framework where users of the library can create Controller classes. Controllers can have their methods decorated with @scan(0.1) to create update loops called at the given period by the control system. When a Controller is instantiated and passed to the control system, bound methods will be created to allow calling them with an implicit self parameter outside the class context.

The decorator needs to annotate the methods with a Scan object so that during initialisation the attributes of the controller can be searched for these annotated methods. The methods must remain callable on the controller instance directly. I have tried various things, but not managed to do this in a type-safe way that satisfies pyright.

Some options I have explored:

1. Add Scan as an attribute to the function in the scan wrapper and search for attributes that satisfy a protocol with an attribute method. This clearly isn't type-safe because the function does not have a method attribute.

2. Replace method with Scan directly and define __call__ with the same signature as the method. This doesn't work because when the method is replaced, accessing the callable class via controller.do_update() does not implicitly pass self as the first argument any more, so the method would have to be called like controller.do_update(controller).

from typing import Callable, Coroutine


ScanCallback = Callable[["Controller"], Coroutine[None, None, None]]

class Scan:
    def __init__(self, fn: ScanCallback, period: float) -> None:
        self._fn = fn
        self._period = period
    
    async def __call__(self, controller: "Controller") -> Coroutine[None, None, None]:
        return self._fn(controller)


def scan(period: float) -> Callable[[ScanCallback], Scan]:
    def wrapper(fn: ScanCallback) -> Scan:
        return Scan(fn, period)

    return wrapper


class Controller:

    @scan(1)
    async def do_update(self):
        print("Do update")


c = Controller()

coro = c.do_update(c)

Playground

3. 2, with the addition of replacing the callable Scan class again in Controller.__init__ with a BoundScan (which is currently only used to call these methods outside of the class context) whose __call__ manually inserts the Controller instance as the first argument. I don't think it is possible to do this in a way that pyright can verify statically, and it is pretty horrible.

I played around with ParamSpec and Concatenate based on this, but I don't think that helps here because it is returning a callable object, not a new function, and I don't think __call__ can be type hinted in the same way. I also looked how fastapi implements its route decorators by storing the state in a separate container and not modifying the wrapped methods at all (e.g. app = FastAPI() ... @app.get), but I would like to avoid this and use decorator functions only if possible.

Does anyone have any suggestions or examples of how to do this? Or, if this just isn't a sensible thing to do, a different approach that would produce the same API?

[rayansostenes]

I played around a little bit with ParamSpec and Concatenate and was able to make it work, I had to use a BoundedScan class, but IMO the DX for writing the Controller classes are pretty good.

from collections.abc import Callable
from typing import Concatenate
from typing import Self
from typing import overload
from typing import reveal_type


class BoundedScan[**P, TController: Controller, Result]:
    def __init__(
        self,
        fn: Callable[Concatenate[TController, P], Result],
        period: float,
        controller: TController,
    ) -> None:
        self._fn = fn
        self._period = period
        self._controller = controller
        # Probably setup the timer logic here
        # ...

    def __call__(self, *args: P.args, **kwds: P.kwargs) -> Result:
        return self._fn(self._controller, *args, **kwds)


class Scan[**P, TController: Controller, Result]:
    def __init__(
        self,
        fn: Callable[Concatenate[TController, P], Result],
        period: float,
    ) -> None:
        self._fn = fn
        self._period = period

    @overload
    def __get__(self, instance: None, owner: type[TController]) -> Self: ...
    @overload
    def __get__(self, instance: TController, owner: type[TController]) -> BoundedScan[P, TController, Result]: ...

    def __get__(self, instance: TController | None, owner: type[TController]) -> BoundedScan[P, TController, Result] | Self:
        if instance is None:
            return self
        return BoundedScan(self._fn, self._period, instance)


def scan[**P, TController: Controller, Result](
    period: float,
) -> Callable[[Callable[Concatenate[TController, P], Result]], Scan[P, TController, Result]]:
    def wrapper(fn):
        return Scan(fn, period)

    return wrapper


class Controller:
    @scan(1)
    async def do_update(self, x: int, y: str) -> int: ...


async def main():
    c = Controller()
    reveal_type(c.do_update)
    #           ^^^^^^^^^^^
    # Type of "c.do_update" is "BoundedScan[(x: int, y: str), Controller, Coroutine[Any, Any, int]]"

    reveal_type(await c.do_update(1, "2"))  # Type of "await c.do_update(1, "2")" is "int"

    await c.do_update("1", 2)
    #                 ^^^^^^
    #                 Argument of type "Literal['1']" cannot be assigned to parameter "x" of type "int" in function "__call__"
    #                   "Literal['1']" is not assignable to "int"
    #
    #                 Argument of type "Literal[2]" cannot be assigned to parameter "y" of type "str" in function "__call__"
    #                   "Literal[2]" is not assignable to "str"

[GDYendell]

Thanks @rayansostenes ! This looks interesting and I will give it a try when I get some time.