drs
drs

Reputation: 5797

How can I extend a library's decorator?

I would like to extend a library's decorator. I know that I can just call both decorators:

@my_decorator
@lib_decorator
def func():
    pass

But I would like to avoid having to pass @lib_decorator to each function each time. I would like my decorator to automatically decorate func() with lib_decorator. How can I do this? Can they be nested?

Upvotes: 15

Views: 5981

Answers (2)

drs
drs

Reputation: 5797

You can incorporate the lib's decorator within yours. For simple, argument-less decorators, it's rather straight-forward:

def my_decorator():
    @lib_decorator  # <--- Just include the lib's decorator here
    def inner:
        func(*args, **kwargs)
    return inner

It's a bit trickier for decorators that have arguments. Just remember that your decorator is replacing the decorated function with the inner-most function. So that's the one you need to decorate. So if you call your decorator with args, e.g.

@my_decorator(arg)
def func():
    pass

Then decorate the inner function with the lib decorator:

def my_decorator(arg):
    def wrapper(func):

        @lib_decorator  # <--- Just include the lib's decorator here
        def inner(*args, **kwargs):
            func(*args, **kwargs)

        return inner

    return wrapper

Or, using the class form of the decorator function:

class my_decorator():
    def __init__(self, arg):
        pass

    def __call__(self, func):
        @lib_decorator  # <--- Just include the lib's decorator here
        def inner(*args, **kwargs):
            func(*args, **kwargs)

        return inner

Upvotes: 17

shx2
shx2

Reputation: 64338

You can easily transform a decoration like yours:

@my_decorator
@lib_decorator
def func():
    pass

To this simpler decoration, using function composition:

my_composed_decorator = lambda func: my_decorator(lib_decorator(func))

@my_composed_decorator
def func():
    pass

Upvotes: 4

Related Questions