I have a Flask app with a custom decorator:
# from functools import wraps, partial
# import flask as f
def require_login(endpoint=None, needs_admin=False):
if endpoint is None:
return partial(require_login, needs_admin=needs_admin)
@wraps(endpoint)
def wrapper(*args, **kwargs):
if "user" not in f.session:
return "Not logged in!", 401
elif needs_admin and not get_current_user()["admin"]:
return "Not an admin!", 403
else:
return endpoint(*args, **kwargs)
return wrapper
It’s used like this:
@require_login
@app.route("/protected")
def protected():
return "Logged in as user"
@require_login(needs_admin=True)
@app.route("/admin")
def admin():
return "Logged in as admin"
These both call get_current_user(), which essentially just returns an object from f.session["user"], which may have "admin" set to True or False.
The intent is that @require_login checks to ensure that the user exists and is sufficiently privileged before allowing the wrapped function to run.
This seems to work fine with authenticated users. However, upon trying to load a resource without sufficient privilege:
Traceback (most recent call last):
# ...
File "users.py", line 73, in protected
return f"Logged in as {get_current_user()}"
^^^^^^^^^^^^^^^^^^^^ # The function is still called!
File "users.py", line 10, in get_current_user
return f.current_app.um.get(f.session['user'])
^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/site-packages/flask/sessions.py", line 80, in __getitem__
return super().__getitem__(key)
^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'user'
>Solution :
The function that is being added as a Flask endpoint is the undecorated function.
Just switch the order of your decorators so that your wrapper is actuallly seem by flask:
@app.route("/admin")
@require_login(needs_admin=True)
def admin():
return "Logged in as admin"
Decorators are passed the function object, and usually it is that that is worked upon. Flask’s .route decorator could operate different, and, for example, take note of the function name, and them, retrieve that name from the module where it is defined when the route is matched. But that would make little sense, and possibly be an open door for even more subtle bugs.
Instead, it just behaves as expected from a decorator: the decorated callable is annotated in a registry as the target for the specified route. What it does differently from most decorators is that, as it does not need any custom code to wrap the function itself, it returns the same function object. If you apply your require_login after that decorator, however, it will mark the original function as the entry point, not your wrapper.
As for the order of decorators application: further from the function is applied last – it may seen strange for some people, but natural for others. Just think of decorators as function calls, and it may make more sense: you are calling your decorator on whatever is returned by @app.route, after it ran.