Combine two python decorators into one

问题:

Here are two decorators I'd like to combine as they are pretty similar, the difference is how a not authenticated user is handled. I'd prefer to have one single decorator that I can call with an argument.

# Authentication decorator for routes
# Will redirect to the login page if not authenticated
def requireAuthentication(fn):
    def decorator(**kwargs):
        # Is user logged on?
        if "user" in request.session:
            return fn(**kwargs)
        # No, redirect to login page
        else:
            redirect('/login?url={0}{1}'.format(request.path, ("?" + request.query_string if request.query_string else '')))
    return decorator

# Authentication decorator for routes
# Will return an error message (in JSON) if not authenticated
def requireAuthenticationJSON(fn):
    def decorator(**kwargs):
        # Is user logged on?
        if "user" in request.session:
            return fn(**kwargs)
        # No, return error
        else:
            return {
                "exception": "NotAuthorized",
                "error" : "You are not authorized, please log on"
            }
    return decorator

And currently I'm using these decorators for specific routes, e.g.

@get('/day/')
@helpers.requireAuthentication
def day():
    ...

@get('/night/')
@helpers.requireAuthenticationJSON
def night():
    ...

I'd much more prefer this:

@get('/day/')
@helpers.requireAuthentication()
def day():
    ...

@get('/night/')
@helpers.requireAuthentication(json = True)
def night():
    ...

I'm on python 3.3 using the Bottle framework. Is it possible to do what I want? How?

回答1:

Just add another wrapper to capture the json parameter:

def requireAuthentication(json=False):
    def decorator(fn):
        def wrapper(**kwargs):
            # Is user logged on?
            if "user" in request.session:
                return fn(**kwargs)

            # No, return error
            if json:
                return {
                    "exception": "NotAuthorized",
                    "error" : "You are not authorized, please log on"
                }
            redirect('/login?url={0}{1}'.format(request.path, ("?" + request.query_string if request.query_string else '')))
        return wrapper
    return decorator

I've renamed your original requireAuthentication function to decorator (because that is what that function did, it decorated fn) and renamed the old decorator to wrapper, the usual convention.

Whatever you put after the @ is an expression, evaluated first to find the actual decorator function. @helpers.requireAuthentication() means you want to call requireAuthentication and it's return value is then used as the actual decorator for the function the @ line applies to.

回答2:

You can create wrapper for both of these decorators:

def requireAuthentication(json=False):
    if json:
        return helpers.requireAuthenticationJSON
    else:
        return helpers.requireAuthentication

Or

import functools
# Authentication decorator for routes
# Will redirect to the login page if not authenticated
def requireAuthentication(json=False):
    def requireAuthentication(fn):
        @functools.wraps(fn)
        def decorator(*args, **kwargs):
            # Is user logged on?
            if "user" in request.session:
                return fn(*args, **kwargs)
            if json:
                 return {
                "exception": "NotAuthorized",
                "error" : "You are not authorized, please log on"
            }
            return redirect('/login?url={0}{1}'.format(request.path, 
                                                       ("?" + request.query_string if request.query_string else '')))
        return decorator
    return requireAuthentication