Home Articles Categories Series
Pythonise Just now

Custom Flask decorators | Learning Flask Ep. 28

Using Python decorators to add another layer of functionality to Flask routes


Article Posted on by in Flask
Julian Nash · 7 months ago in Flask

Just like decorators in Python allow us to add additional functionality to functions. View decorators in Flask allow us to add additional functionality to routes.

As a view (or route) in Flask is a Python function, we can implement our own decorators to add another layer of functionality to them.

Python decorators

Here's an example of a basic Python decorator:

from functools import wraps

# Defining our custom decorator
def my_decorator(function):
    @wraps(function)
    def wrapper(a, b, c):
        print("wrapper running!")
        a += 1
        b += 2
        c += 3
        return function(a, b, c)
    return wrapper

# Using it to decorate a function
@my_decorator
def my_function(a, b, c):
    print("my_function running!")
    print(a, b, c)

my_function(a=1, b=2, c=3)

Running the above prints:

wrapper running!
my_function running!
2 4 6

Pay attention to the order of execution.

  • The first call to print() was from within our decorator
  • Followed by the 2 calls to print() in my_function

You'll also notice, the decorator changed the values we passed into the my_function call.

By decorating a function with @my_decorator, the function directly below it is passed into the my_decorator code as the function agument.

The original my_function is replaced with the function we returned in our decorator.

You'll also notice we've imported wraps from functools.

@wraps is not required, but helps us by copying the function docsting, name, attributes etc. from the original function to the copy of the function inside the decorator!

Here's another example that accepts arguments in the decorator:

def html_tag_generator(tag, attrs):
    def decorator(function):
        @wraps(function)
        def wrapper(text):
            attr_string = " ".join(f"{k}='{v}'" for k, v in attrs.items())
            text = f"<{tag} {attr_string}>{text}</{tag}>"
            return function(text)
        return wrapper
    return decorator


@html_tag_generator(tag="div", attrs={"class": "container col-s-12", "id": "mydiv"})
def modify_text(text):  # <- `text` has been modified by the decorator!
    print(text)


modify_text("Python")

running this prints:

<div class='container col-s-12' id='mydiv'>Python</div>

Flask decorators

If you've used Flask before, you'll be very familiar with many of Flask's decorators, such as:

@app.route

Use the @app.route decorator to define the routes in your application:

@app.route("/profile", methods=["GET", "POST"])
def profile:
    return render_template("profile.html")
@app.before_request

Use the @app.before_request decorator to trigger a function to run before every request:

@app.before_request
def do_before_request:
    connect(**app.config["MONGO_CONNECTION"])
    g.conn = psycopg2.connect(**app.config["POSTGRES_CONNECTION"])
    g.user = get_user_from_session()
@app.errorhandler(err_code)

Use the @app.errorhandler decorator to catch errors:

@app.errorhandler(404)
def page_not_found(e):
    return render_template("error_pages/404.html"), 404

As you can see, decorators are ubiquitous in Flask!

Custom Flask decorators

Armed with the knowledge to write your own custom decorators, here's an example:

def superuser(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):

        if not g.user.superuser:
            flash("You do not have permission to view that page", "warning")
            abort(404)

        return f(*args, **kwargs)

    return decorated_function

Let's step through each line of our decorator:

def superuser(f):

Here we define a function, which will be the name of our decorator. The f argument is the function which we'll decorate.

@wraps(f)

We'll use the @wraps decorator and pass it the function (f), copying the original function f's metadata to the new function we're about to define.

def decorated_function(*args, **kwargs):

The name of this function doesn't matter, however it's used to modify the behaviour and values passed into the original function.

if not g.user.superuser:
    flash("You do not have permission to view that page", "warning")
    abort(404)

Here, we're checking the g.user has a True value for the superuser attribute. If not, we're calling abort(404).

return f(*args, **kwargs)

Returns the newly modified function and any arguments passed into it. You'll notice that we haven't modified any of the *args or **kwargs passed into the function, but we still need to return them.

return decorated_function

This line just returns the new function to the parent function so it can be returned.

Using the decorator

Now that we have our decorator, we can use it:

@app.route("/users")
@superuser
def users():
    user = g.user
    return render_template("users/users.html", user=user)

Decorators can be stacked, and you'll notice we've used our @superuser just under the @app.route decorator.

Any time a request is sent to this route, it will trigger the @superuser decorator to run. If someone tried to access the /users URL and doesn't have the right permissions defined in our decorator, the request will be aborted.

Decorator arguments

Sometimes, it's useful to pass arguments into our custom decorators:

def restricted(access_level):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            print(access_level)
            return func(*args, **kwargs)
        return wrapper
    return decorator

We now have access to the access_level value passed into our decorator.

Usage:

@app.route("/dashboard")
@restricted(access_level="admin")
def dashboard():
    user = g.user
    return render_template("dashboard/dashboard.html", user=user)

This can be useful for reusing a decorator and passing it different values for different functionality:

def restricted(access_level):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not g.user.access_level == access_level
                abort(403)
            return func(*args, **kwargs)
        return wrapper
    return decorator

You may want to decorate views, passing in different values to the decoartor and letting it handle the validation:

@app.route("/admin")
@restricted(access_level="admin")
def admin_zone():
    user = g.user
    return render_template("admin/dashboard.html", user=user)

@app.route("/superuser")
@restricted(access_level="superuser")
def superuser_zone():
    user = g.user
    return render_template("superuser/dashboard.html", user=user)

@app.route("/user")
@restricted(access_level="user")
def user_zone():
    user = g.user
    return render_template("user/dashboard.html", user=user)

Wrapping up

Decorators in Flask are a great way to add an additional layer of functionality to a route and provide a nice way to keep your code DRY.

Last modified · 16 Apr 2019
Did you find this article useful?
This work is licensed under a Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License
Contents
Loading...