How to make decorator work with async function?

Viewed 1075

Generic question:

I have a decorator @some_package.decorator that expects the function to wrap to return type T.

### runs
@some_package.decorator
def my_func(param) -> int:
    return len(param)

### I also want to be able to wrap an async function
@some_package.decorator
async def my_func(param):
    await asyncio.sleep(1)
    return len(param)

However the function can either be sync or async and so may return coroutine[Any].

How can I make the decorator await the function

My specific use case: I have some code that works well when sync=True

import uplink
class MyAPI(Consumer):
        def __init__(
        self,
        password: str,
        sync: bool = True,
        auto_auth: bool = True,
        **kwargs
    ):
        self.sync = sync
        self.client = None

        if not self.sync:
            self.client = AiohttpClient()

        # initialise the super class
        super(GLinet, self).__init__(client=self.client, **kwargs)

        if auto_auth:
                self._login(password)

    @uplink.returns.json(key="token")
    @uplink.post("router/login")
    def _login(self, pwd: uplink.Field):
        """fetches token"""

However when sync=False the API I am using makes and _login() and async function. As such the decorator uplink.returns.json becomes upset as it expects the function it wraps to prove a response that is can parse but instead it receives a co routine.

How can I make the uplink.returns.json accept an async function and await a return value that it can parse.

I tried something like this

def dec(fn):
    if asyncio.iscoroutinefunction(fn):
        @returns.json
        async def wrapper(*args, **kwargs):
            print("wrapping async function")
            print(fn)
            return await fn()
        return wrapper
    else:
        @returns.json
        def wrapper(*args, **kwargs):
            print("wrapping sync function")
            print(fn)
            return fn()
        return wrapper

Some people have developed their own decorators that can take both tpyes of functions, but I haven't found anyone who has modified another decorator

def dec(fn):
    if asyncio.iscoroutinefunction(fn):
        @wraps(fn)
        async def wrapper(*args, **kwargs):
            print(fn, args, kwargs)  # <function foo at 0x10952d598> () {}
            await asyncio.sleep(5)
            print("done with wrapper, going to call fn")
            return await fn()
        return wrapper
    else:
        @wraps(fn)
        def wrapper(*args, **kwargs):
            print(fn, args, kwargs)  # <function bar at 0x108fb5a60> () {}
            time.sleep(5)
            print("done with wrapper, going to call fn")
            return fn()
        return wrapper
0 Answers
Related