We know that when using decorators in Python you should always use functools.wraps/update_wrapper on the function returned by the decorator. Apart from setting the __name__, __doc__, __module__... attributes of the new/wrapper function with those of the original one, it also adds a __wrapped__ attribute that points to the original function. What it does not do is adding information to the function about the decorator that has been applied. So while we have a way to refer to the original function via __wrapped__, we can not check if a decorator has been applied to the function.
Obviously our decorator could just add a __decorator__ attribute to the wrapper function that it returns, but well, we have to repeat that logic in each of our decorators, and can not do anything with already existing decorators. So the nice way to do this would be having a function (let's call it empower()) that we can apply to an existing decorator, obtaining a new decorator that applies the original decorator and then sets the __decorator__ attribute in the decorated function. This empower function is a decorator factory (receives a decorator and creates a new decorator) and indeed can be applied (at least in some cases, we'll see it later) as a decorator itself when defining the initial decorator, so it could be seen as a sort of meta-decorator (a decorator that decorates and creates decorators).
A function can be decorated by multiple decorators, so shouldn't we better have a __decorators__ attribute with that list of decorators? Well, the __wrapped__ attribute points to the function being decorated in this step (so functools.wraps does not check if the function being decorated has in turn a __wrapped__ attributes). So if we have a chain of decorators we'll have to traverse a chain of __wrapped__ attributes to get to the source function. I mean:
"""
Veryfying that if multiple decorators are applied, functools.__wrapped__ points to the previous decorated function in the chain, not directly to the original function.
"""
from functools import wraps
def start_call(func):
@wraps(func)
def wrapper(*args, **kwargs):
print(f"Starting function: {func.__name__}")
return func(*args, **kwargs)
return wrapper
def end_call(func):
@wraps(func)
def wrapper(*args, **kwargs):
result = func(*args, **kwargs)
print(f"Ending function: {func.__name__}")
return result
return wrapper
@start_call
@end_call
def do_something(a, b):
return a + b
# Example usage
if __name__ == "__main__":
# do_something has 2 levels of decoration
print(f"Result: {do_something(5, 10)}")
print("---------------- Unwrapping decorators ----------------")
unwrapped1 = do_something.__wrapped__
print(f"Result: {unwrapped1(5, 10)}")
print("----------------")
unwrapped2 = do_something.__wrapped__.__wrapped__
print(f"Result: {unwrapped2(5, 10)}")
# Starting function: do_something
# Ending function: do_something
# Result: 15
# ---------------- Unwrapping decorators ----------------
# Ending function: do_something
# Result: 15
# ----------------
# Result: 15
So that's also the approach I've followed here. I add a __decorator__ attribute to each decorated function, and have a get_decorators helper function that will traverse that __decorator__ chain to get all the decorators. My empower decorator provides an additional functionality, if the decorator being decorated does not apply wraps() to the original function, it does it. Let's see an implementation of this empower decorator and the associated get_decorators() function.
def empower(decorator):
def empowered_decorator(func):
decorated_fn = decorator(func)
decorated_fn.__decorator__ = decorator
# if the original decorator has not used wraps, we add it here
if not hasattr(decorated_fn, '__wrapped__') or decorated_fn.__wrapped__ != func:
wraps(func)(decorated_fn)
return decorated_fn
return empowered_decorator
def get_decorators(func):
decorators = []
while cur_decor := getattr(func, '__decorator__', None):
decorators.append(func.__decorator__)
func = func.__wrapped__
return decorators
Given the previously defined start_call and end_call decorators, we can empower them at the time they are applied to a function, like this:
@(empower(start_call(">>>")))
@(empower(end_call))
def do_something2(a, b):
return a + b
print(do_something2(7, 3))
print(f"decorators applied to do_something2: {[dec.__name__ for dec in get_decorators(do_something2)]}")
# decorators applied to do_something2: ['intermediate', 'end_call']
Being used at the time a decorator is being applied, rather that at the time when a decorator is defined, the empower decorator works naturally both for decorators that expect parameters and decorators that do not (other than the function being decorated). A decorator that expects parameters does indeed create a new decorator that traps the provided parameters in its closure for being then invoked with the function to be decorated, so in both cases empower ends up receiving a decorator that just expects a function.
We can also apply it when a decorator is being defined, but only for decorators that do not expect parameters (other than the function being decorated itself). For applying it at definition time to decorators that expect parameters, we need a different implementation, that I've called empower_dec_with_params. So all in all we have:
# intended to be used when defining a decorator with parameters
def empower_dec_with_params(decorator):
def outer_decorator(*args, **kwargs):
def inner_decorator(func):
decorated_fn = decorator(*args, **kwargs)(func)
decorated_fn.__decorator__ = decorator
# if the original decorator has not used wraps, we add it here
if not hasattr(decorated_fn, '__wrapped__') or decorated_fn.__wrapped__ != func:
wraps(func)(decorated_fn)
return decorated_fn
return inner_decorator
return outer_decorator
@empower_dec_with_params
def start_call(prepend: str = ""):
def intermediate(fn):
@wraps(fn)
def wrapper(*args, **kwargs):
print(f"{prepend} Starting function: {fn.__name__}")
return fn(*args, **kwargs)
return wrapper
return intermediate
@empower
def end_call(func):
@wraps(func)
def wrapper(*args, **kwargs):
result = func(*args, **kwargs)
print(f"Ending function: {func.__name__}")
return result
return wrapper
@start_call(">>>")
@end_call
def do_something(a, b):
return a + b
print(do_something(5, 10))
print(f"decorators applied to do_something: {[dec.__name__ for dec in get_decorators(do_something)]}")
# >>> Starting function: do_something
# Ending function: do_something
# 15
# decorators applied to do_something: ['start_call', 'end_call']
No comments:
Post a Comment