pythonattributesdecoratorpython-decoratorsfunction-attributes

Decorator to set attribute of function


I want different functions to be executable only if the logged in user has the required permission level.

To make my life more simple I want to use decorators. Below I attempt to set attribute permission on 'decorated' functions - as shown below.

def permission(permission_required):
    def wrapper(func):
        def inner(*args, **kwargs):
            setattr(func, 'permission_required', permission_required)
            return func(*args, **kwargs)
        return inner
    return wrapper

@permission('user')
def do_x(arg1, arg2):

    ...

@permission('admin')
def do_y(arg1, arg2):
    ...

But when I do:

fn = do_x
if logged_in_user.access_level == fn.permission_required:
    ...

I get an error 'function' object has no attribute 'permission_required'

What am I missing?


Solution

  • You are checking the attribute on the inner (wrapper) function, but set it on the original (wrapped) function. But you need a wrapper function at all:

    def permission(permission_required):
        def decorator(func):
            func.permission_required = permission_required
            return func
        return decorator
    

    Your decorator needs to return something that'll replace the original function. The original function itself (with the attribute added) will do fine for that, because all you wanted to do is add an attribute to it.

    If you still need a wrapper, then set the attribute on the wrapper function instead:

    from functools import wraps
    
    def permission(permission_required):
        def decorator(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                # only use a wrapper if you need extra code to be run here
                return func(*args, **kwargs)
            wrapper.permission_required = permission_required
            return wrapper
        return decorator
    

    After all, you are replacing the wrapped function with the wrapper returned by the decorator, so that's the object you'll be looking for the attribute on.

    I also added the @functools.wraps() decorator to the wrapper, which copied across important identifying information and other helpful things from func to the wrapper, making it much easier to work with.