import functools
import asyncio


def task_handled_errors(func):
    """
    Decorator for custom tasks. Can *only* wrap coroutines. Any raised exceptions will call the
    KazTron general error handler.
    """
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        # noinspection PyBroadException
        try:
            return await func(*args, **kwargs)
        except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
            raise
        except Exception:
            await args[0].core.on_error(func.__name__)
    return wrapper