Open bckohan opened 1 week 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
From upstream: https://github.com/fastapi/typer/issues/950
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,
)
Is this possible at the django-typer interface or does it require upstream changes?