Python: how to make this more generic - probably with a decorator?

Question

currently I have this code

@app.route("/protect1")
def protect1():
    if not session.get('logged_in'):
        session['next'] = "/protect1"
        return render_template('login.html')
    else:
        return "This is the first protected page protect1"   

@app.route("/protect2")
def protect2():
    if not session.get('logged_in'):
        session['next'] = "/protect2"
        return render_template('login.html')
    else:
        return "This is the second protected page protect2"   

in my flask app, everything working fine. Only it is not nice that I will need to repeat for each function (view) the if/else combination.

I would prefer to have some generic way, like this pseude-code:

@checklogin
@app.route("/protect1")
def protect1():
    return "This is the first protected page protect1"

@checklogin
@app.route("/protect2")
def protect2():
    return "This is the second protected page protect2"

One challenge here is that the @checklogin decorator would need to know the app.route path (eg "/protect1") in order to be able to set session['next'] correctly. I have no idea how to pass this parameter to the decorator, especially how to find it out in the first place. In other words, how does the function protect1() know that it is decorated with @app.route and which parameter ("/protect1") has been passed to that app.route decorator?

Answer 1

The decorator can look up the path on request ; either using the URL that was loaded (available as request.url ) or the request.endpoint attribute :

from functools import wraps
from flask import request, session

def checklogin(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if not session.get('logged_in'):
            session['next'] = request.url
            return render_template('login.html')
        return f(*args, **kwargs)
    return wrapper

Do place the decorator after the app.route() decorator or it'll not be registered as the handler for the route:

@app.route("/protect1")
@checklogin
def protect1():
    return "This is the first protected page protect1"