click icon indicating copy to clipboard operation
click copied to clipboard

[Feature] embedding Click inside asyncio-based apps

Open achimnol opened this issue 3 years ago • 1 comments

#2033 and #85 addresses how to run "self-contained" async functions inside Click commands and it is easily achievable by using asyncio.run(). Here "self-contained" means that each function runs to the completion of the entire event loop's lifecycle (i.e., when Click serves a CLI app).

Though, I feel there is lack of support from Click when I try to embed Click inside an asyncio event loop that is already established and invoke coroutines inside the synchronous command functions (e.g., in a telnet server handler of prompt_toolkit), because asyncio.run() and loop.run_until_complete() cannot be executed inside an already-running event loop started by another asyncio.run() and its coroutines/callbacks.

My current solution is:

import asyncio
import functools
import shlex
import traceback
from contextvars import ContextVar, copy_context

import click
from prompt_toolkit import PromptSession

command_done: ContextVar[asyncio.Event] = ContextVar("command_done")


@click.group(add_help_options=False)
def monitor_cli():
    pass


async def cli_loop():
    ...
    prompt_session: PromptSession[str] = PromptSession()
    while True:
        try:
            user_input = (
                await prompt_session.prompt_async("command> ")
            ).strip()
            args = shlex.split(user_input)
        except (EOFError, KeyboardInterrupt, asyncio.CancelledError):
            return
        except Exception:
            _print_error(traceback.format_exc())
        else:
            command_done_event = asyncio.Event()
            command_done_token = command_done.set(command_done_event)
            try:
                ctx = copy_context()
                ctx.run(
                    monitor_cli.main,
                    args,
                    prog_name="",
                    standalone_mode=False,  # type: ignore
                )
                # Before displaying a new prompt, we need to wait until the command is done.
                await command_done_event.wait()
            finally:
                command_done.reset(command_done_token)


if __name__ == "__main__":
    asyncio.run(cli_loop())


...


def custom_help_option(cmdfunc):
    """
    A custom help option to ensure setting `command_done_event`.
    If we just rely on the Click's vanilla help option,
    `await command_done_event.wait()` inside `cli_loop()` will never return.
    """

    @auto_command_done  # <-- this is what we need to extend here
    def show_help(ctx: click.Context, param: click.Parameter, value: bool) -> None:
        if not value:
            return
        click.echo(ctx.get_help(), color=ctx.color)
        ctx.exit()

    return click.option(
        "--help",
        is_flag=True,
        expose_value=False,
        is_eager=True,
        callback=show_help,
        help="Show the help message",
    )(cmdfunc)


def auto_command_done(cmdfunc):
    @functools.wraps(cmdfunc)
    def _inner(ctx: click.Context, *args, **kwargs):
        command_done_event = command_done.get()
        try:
            return cmdfunc(ctx, *args, **kwargs)
        finally:
            command_done_event.set()

    return _inner


def auto_async_command_done(cmdfunc):
    @functools.wraps(cmdfunc)
    async def _inner(ctx: click.Context, *args, **kwargs):
        command_done_event = command_done.get()
        try:
            return await cmdfunc(ctx, *args, **kwargs)
        finally:
            command_done_event.set()

    return _inner


@monitor_cli.command()
@custom_help_option
@auto_command_done  # shortcut for normal sync command functions
def do_sync_work():
    ...


@monitor_cli.command()
@custom_help_option
def do_async_work():

    @auto_async_command_done  # shortcut for async command functions
    async def _do_async_work():
        await asyncio.sleep(1)
        ...

    asyncio.create_task(_do_async_work())

The above snippet works as expected, but it brings a quite much amount of boilerplates.

My suggestions are to:

  • Allow adding decorators to the intrinsic help option's callback
  • Add intrinsic decorators like auto_async_command_done and auto_async_command_done with a wrapper API of command_done_event.wait()

Probably we could write an AsyncCommand class (exposed as click.async_command()) and BaseCommand.async_main() to handle these seamlessly.

achimnol avatar Sep 24 '22 08:09 achimnol

The easy way out is to Install asyncclick. Just import asyncclick as click and you can use async commands, callbacks and so on as you please.

Async coverage should be reasonably complete; if you hit a callback that you need to be async which isn't, please file a bug.

smurfix avatar Oct 11 '22 14:10 smurfix