import inspect


def wrap_generator(func):
    """
    Decorator to convert a generator function to an async function which collects
    and returns generator results, returning a list if there are multiple results
    """

    async def _wrapped(*a, **k):
        r, ret = None, []
        gen = func(*a, **k)
        while True:
            try:
                item = gen.send(r)
            except StopIteration:
                break
            if inspect.isawaitable(item):
                r = await item
            else:
                r = item
            ret.append(r)

        if len(ret) == 1:
            return ret.pop()
        return ret

    return _wrapped