Source code for caspia.toolbox.cli.asynccmd

import asyncio
import functools
import signal
import sys

import click
from aiopollen import errors


def async_command(f=None):
    def decorator(f):
        @functools.wraps(f)
        def wrapped(*args, **kwargs):
            try:

                async def main():
                    loop = asyncio.get_event_loop()

                    def signal_handler():
                        raise KeyboardInterrupt()

                    loop.add_signal_handler(signal.SIGTERM, signal_handler)
                    return await f(*args, **kwargs)

                return asyncio.run(main())
            except (asyncio.TimeoutError, errors.TimeoutError):
                click.secho('Timed out.', fg='red', err=True)
                sys.exit(1)
            except KeyboardInterrupt:
                pass

        return wrapped

    if f is None:
        return decorator
    else:
        return decorator(f)