django-typer icon indicating copy to clipboard operation
django-typer copied to clipboard

Support async commands

Open bckohan opened this issue 1 year ago • 5 comments

  • start event loops automatically
  • allow asynchronous execution of chained subcommands

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

bckohan avatar Nov 15 '24 21:11 bckohan

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 avatar Nov 17 '24 18:11 bckohan

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

bckohan avatar Nov 21 '24 21:11 bckohan

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,
        )

pySilver avatar Nov 22 '24 23:11 pySilver

This should really be in click.

Upstream issues:

  1. Great explanation at the click level
  2. Typer issue.

bckohan avatar Feb 13 '25 05:02 bckohan

This is a much thornier issue than It would seem. I'm punting this to a future release - it should best be handled upstream in click.

bckohan avatar Feb 14 '25 07:02 bckohan