async_maybe.py

import asyncio
from functools import wraps
import os

import subprocess
import time

class await_maybe:

    def __init__(self, val):
        self.val = val

def running_async():
    if os.environ.get('RUN_ASYNC'):
        return True
    return False

def sync_switch(sync_func, async_func, *args, **kwargs):
    if running_async():
        func = async_func
    else:
        func = sync_func
    return await_maybe(func(*args, **kwargs))

def async_maybe(func):
    
    @wraps(func)
    async def async_wrapper(*args, **kwargs):
        print('async!!!')
        gen = func(*args, **kwargs)
        if not hasattr(gen, 'send'):
            raise StopIteration(gen)
        val = gen.send(None)
        try:
            while True:
                if not isinstance(val, await_maybe):
                    raise Exception('async_maybe does not work with generators.')
                else:
                    val = gen.send(await val.val)
        except StopIteration as ret_exc:
            return ret_exc.value

    @wraps(func)
    def sync_wrapper(*args, **kwargs):
        print('sync!!!')
        gen = func(*args, **kwargs)
        if not hasattr(gen, 'send'):
            return gen
        val = gen.send(None)
        try:
            while True:
                if not isinstance(val, await_maybe):
                    raise Exception('async_maybe does not work with generators.')
                else:
                    val = gen.send(val.val)
        except StopIteration as ret_exc:
            return ret_exc.value

    if running_async():
        return async_wrapper
    else:
        return sync_wrapper


class CommandRunner:

    def __init__(self, cmd, times):
        self.cmd = cmd
        self.times = times
    
    @async_maybe
    def run_fast(self):
        val = yield sync_switch(self.run_sync, self.run_async)
        yield await_maybe(self.print_summary())
        return val
    
    @async_maybe
    def print_summary(self):
        print('---------')
        yield sync_switch(time.sleep, asyncio.sleep, 5)
        print('Executed "{}" {} times!'.format(self.cmd, self.times))
    
    def run_sync(self):
        return [subprocess.run(self.cmd) for i in range(self.times)]
    
    async def run_async(self):
        processes = [await asyncio.create_subprocess_exec(*self.cmd) for i in range(self.times)]
        return await asyncio.gather(*(process.wait() for process in processes))

runner = CommandRunner(['./thing.sh'], 10)
if running_async():
    print(asyncio.get_event_loop().run_until_complete(runner.run_fast()))
else:
    print(runner.run_fast())