django-commons / django-typer

Use Typer (type hints) to define the interface for your Django management commands.
https://django-typer.rtfd.io
MIT License
88 stars 2 forks source link

Support async callbacks #141

Open bckohan opened 1 week ago

bckohan commented 1 week ago

Is this possible at the django-typer interface or does it require upstream changes?

bckohan commented 5 days ago

To be clear, it'd be nice to be able to do something like:

class Command(TyperCommand, chain=True):

    @command()
    async def sub1(self):
        ...

    @command()
    async def sub2(self):
        ...

Then run both asynchronously like this:

./manage.py command sub1 sub2
bckohan commented 1 day ago

From upstream: https://github.com/fastapi/typer/issues/950

pySilver commented 18 hours ago

That would be awesome. At the moment I'm doing something similar manually.

Here are the utils for commands:

import asyncio
import signal
from collections.abc import Awaitable, Callable
from functools import wraps
from typing import Any, Concatenate, ParamSpec, TypeVar

from django_typer.management import TyperCommand

from config.context import StateManager

T = TypeVar("T", bound="StateManagerCommand")
P = ParamSpec("P")
R = TypeVar("R")

class StateManagerCommand(TyperCommand):
    state_manager: StateManager

    async def startup(self) -> None:
        self.state_manager = StateManager()
        await self.state_manager.startup()

    async def shutdown(self) -> None:
        await self.state_manager.shutdown()

    @staticmethod
    def with_state_manager() -> (
        Callable[
            [Callable[Concatenate[T, P], Awaitable[R]]],
            Callable[Concatenate[T, P], Awaitable[R]],
        ]
    ):
        def decorator(
            func: Callable[Concatenate[T, P], Awaitable[R]],
        ) -> Callable[Concatenate[T, P], Awaitable[R]]:
            @wraps(func)
            async def wrapper(
                self: T,
                *args: P.args,
                **kwargs: P.kwargs,
            ) -> R:
                await self.startup()
                try:
                    return await func(self, *args, **kwargs)
                finally:
                    await self.shutdown()

            return wrapper

        return decorator

def run_in_loop(
    signals: tuple[signal.Signals, ...] = (
        signal.SIGHUP,
        signal.SIGTERM,
        signal.SIGINT,
    ),
    shutdown_func: Callable[[signal.Signals, asyncio.AbstractEventLoop], Any]
    | None = None,
) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, R]]:
    """
    Decorator function that allows defining coroutines with click.

    Args:
        signals: Tuple of signal types to handle
        shutdown_func: Optional callback function for signal handling

    Returns:
        A wrapped coroutine function that handles signal management
    """

    def decorator(
        func: Callable[P, Awaitable[R]],
    ) -> Callable[P, R]:
        @wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            loop: asyncio.AbstractEventLoop = asyncio.get_event_loop()
            if shutdown_func:
                for ss in signals:
                    loop.add_signal_handler(ss, shutdown_func, ss, loop)
            return loop.run_until_complete(func(*args, **kwargs))

        return wrapper

    return decorator

and there is some basic usage:

import typer
from django_typer.management import command

from project.core.management.utils import StateManagerCommand, run_in_loop
from project.core.tasks import awesome_task

class Command(StateManagerCommand):
    @command()
    @run_in_loop()
    @StateManagerCommand.with_state_manager()
    async def default(self) -> None:
        await awesome_task.kiq()
        self.secho("Awesome task scheduled for execution", fg="green")
        self.secho(
            f"State: {self.state_manager.state.keys()}",
            fg=typer.colors.MAGENTA,
        )