How to Write a Decorator in Python Flask to Check Logged In Status

Testing my @logged_in decorator

If you’re writing a web app, chances are there are that there are endpoints which you want to restrict access to based on if a user is logged in or has certain permissions.

You might have a code snippet like this ALL OVER your project repeated for every routing:

if session.get(“logged_in”):
    return render_template(“private_stuff.html”)
else:
    return redirect(“/login”)

This is not just adding a few lines of code, you have to encapsulate your main controller or view logic inside an if-statement for EVERY endpoint you want to restrict access to.

What if you could do this by using only one line of code per endpoint? Well you can with a decorator.

What’s a Decorator?

Photo by Jess Bailey on Unsplash

The super simple definition of a decorator is it’s something that modifies the behavior of a function. In using Flask, probably you have some familiarity with a couple decorators already such as @app.route(), @classmethod, @staticmethod etc.

To elaborate, a decorator is in itself a function that that takes in another function as an argument and “wraps” around the function passed into it. It took me a bit to wrap my head around this (no pun intended) as it’s easier explained than implemented.

An example of a very simple decorator:

def my_decorator(a_function):
    print(“I decorate thee”)
    return a_function()
 
@my_decorator
def print_hello():
    print(“hello”)

Its output will be:

>>> I decorate thee
>>> hello

But this doesn’t work with functions that take in arguments, so you need an inner function to handle the *args and **kwargs assuming a given function you’re wrapping may have arguments.

def my_decorator(a_function):
    def inner(*args, **kwargs):
        print("I decorate thee")
        return a_function(*args, **kwargs)
    return inner
@my_decorator
def print_color(color):
    print(color)
print_color("Red")

The output will be:

>>> I decorate thee
>>> Red

In the above example, the decorator doesn’t do anything with the arguments passed into print_color()

We can grab the keyword argument for example and write the decorator to change the output accordingly.

def my_decorator(a_function):
    def inner(*args, **kwargs):
        if kwargs.get('color') == "Red":
            print("oh my favorite color!")
        else:
            print("here's a boring color:")
        return a_function(*args, **kwargs)
    return inner
@my_decorator
def print_color(color):
    print(color)
print_color(color="Red")
print_color(color="Blue")

In this case the output will be:

>>> oh my favorite color!
>>> Red
>>> here’s a boring color:
>>> Blue

Note: for this to work, the keyword “color” has to be used in the function calls.

Now this is by no means the best or even a very comprehensive explanation of decorators (aka “wrapper functions”) but we’ve got enough to go on in order write one for our Flask app.

Another detail is the @wraps decorator from functools which is used to preserve the metadata of the wrapped function. It’s not typically necessary but it’ll help us in constructing our decorator for Flask. Here’s the official docs for @wraps .

The @logged_in Decorator

Let’s assume two things; you’ve given Session a Boolean parameter named “logged_in” to determine if the user is logged in or not and that anyone navigating to a private endpoint is redirected to your root endpoint / .

We need to include the following import statement for the @wraps decorator we’re going to use:
from functools import wraps

We can now write the decorator function like so:

def logged_in(f):
    @wraps(f)
    def decorated_func(*args, **kwargs):
        if session.get("logged_in"):
            return f(*args, **kwargs)
        else:
            return redirect("/")
    return decorated_func

The old way and the new way!

Whereas before you’d need to write your protected endpoint like so:

@app.route(“/private_stuff”)
def private_stuff():
    if session.get(“logged_in”):
        return render_template(“private.html”)
    else:
        return redirect(“/”)

You can now write it like this:

@app.route(“/private_stuff”)
@logged_in
def private_stuff():
    return render_template(“private.html”)

You can see even in this very basic example, your route method becomes much simpler to write. In a full stack application with many endpoints in which you’ll have different routing rules, validations and protections, a decorator like this will make it much easier to keep track of.

Since the session object is accessible everywhere, the decorator really just needs to pass the *args and **kwargs through without doing anything with them. But if you wanted to change the routing based on parameters passed to your routing method, there are many ways to do that too as in the above decorator examples.

Ultimately the benefits of creating your own decorators will be to have more streamlined code, fewer opportunities for error and better readability for your routing methods. I’m looking forward to seeing the even better decorators that you’ll come up with!