Higher-order function on async coroutine

Question

Let's say I have a function :

def f(x):
    return {"x": x}

I can create a higher-order function that behaves like so :

def augment(func):
    def augmented_function(x):
         return {**func(x), "metadata": "test"}
    return augmented_function

And then augment(f)(1) will return {"x": 1, "metadata": "test"}.

But what if f is an async coroutine, this augment function does not work (RuntimeWarning: coroutine 'f' was never awaited and TypeError: 'coroutine' object is not a mapping) - I want the augmented function to be a coroutine that can be awaited :

async def f(x):
    return {"x": x}

def augment_async(coro):
   xxx

augment_async(f)(1) # Should return <coroutine object xxx>
await augment_async(f)(1) # Should return {"x": 1, "metadata": "test"}

Does anyone know how to write augment_async in this case?

Thanks.

EDIT : Bonus question.

How to write augment_async such as await augment_async(f(1)) returns {"x": 1, "metadata": "test"}?

Answer 1

It is sufficient to make the inner function async, so that it can await the wrapped function:

def augment_async(func):
    async def augmented_function(x):
         return {**await func(x), "metadata": "test"}
    return augmented_function

await augment_async(f)(1) # evaluates to {"x": 1, "metadata": "test"}

---

To “augment” an instantiated coroutine f(1), a single layer is sufficient:

 async def direct_async(coro):
     return {**await coro, "metadata": "test"}

Note that unlike augment_async, which produces a factory for coroutines, direct_async produces a coroutine directly - its result may be awaited only once.