« Back to Index

[Python Decorator with optional arguments]

View original Gist on GitHub

Tags: #tags: python, decorator

Python Decorator with optional arguments.py

import asyncio
import time

from functools import wraps

def dec(ttl=10):
    print('ttl', ttl)

    def real_dec(fn):
        if asyncio.iscoroutinefunction(fn):
            @wraps(fn)
            async def wrapper(*args, **kwargs):
                print('fn args kwargs', fn, args, kwargs)
                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:', fn, args, kwargs)
                time.sleep(5)
                print("done with wrapper, going to call fn")
                return fn()
            return wrapper

    return real_dec

@dec(ttl=123)
async def foo():
    return await asyncio.sleep(5, result="async function done")

@dec(ttl=456)
def bar():
    time.sleep(5)
    return "sync function done"

loop = asyncio.get_event_loop()
result = loop.run_until_complete(foo('FOO!'))

print(result)
print(bar('BAR!'))

"""
output...

ttl 123
ttl 456

fn args kwargs <function foo at 0x10cfa4620> ('FOO!',) {}
done with wrapper, going to call fn
async function done

fn args kwargs: <function bar at 0x10cc7b8c8> ('BAR!',) {}
done with wrapper, going to call fn
sync function done
"""