Using python decorators to implement guards

One cool technique that I learnt while going through the Django code was using python decorators to implement guards.

What are guards?

Take a look at this bit of pseudo-code

if condition1:
  if condition2:
    {do_something}
  else:
    return error2
else:
  return error1

This is a common pattern where you do something provided condition1 and condition2 are false. The problem with this code is that it is difficult to seperate out the core logic of the function contained in {do_something} and the error handling code in the rest of the function. Another disadvantage is that the condition is at the top of the function, while the failure action is at the bottom. This makes it difficult to correlate the condition with the failure action.

The solution is to refactor the code to use guards.

if not condition1:
  return error1

if not contidion2:
  return error2

{do_something}

Guards are the conditions at the top of the function. They act like security guards — If the condition passes you go through, otherwise you leave the function. It is now a lot easier to see the conditions and the failure actions, and you can easily identify the code logic block by just skipping past the guards.

Python decorators

Python has a decorator feature that allows you to modify the function that it is applied to. Here is an example:

def decorate(fn):
  def _decorate():
    print "before calling"
    fn()
    print "after calling"
  return _decorate

@decorate
def myfunction():
  print "in function"

What we have is a function myfunction that prints the string "in function". To this, we apply the decorator 'decorate' (denoted by the @ symbol). decorate is itself a function that takes one function as a parameter and returns another function. In this case, it takes fn as a parameter and returns _decorate. Everytime myfunction is called, it will actually call the returned function, in this case _decorate. _decorate prints "before calling", then it calls the original function, then prints "after calling". Here is the output

>>> myfunction()
before calling
in function
after calling

Implementing guards using decorators

We can now see how guards can be implemented using decorators. Let me take a real example that I've encountered — my admin page. My tool has an admin page. In order to access this page, you must be logged in, and you must be an admin. If you are not logged in, you need to be redirected to the login page. If you are not an admin, an error message should be displayed. This is how the code would normally have looked

# check for login
try:
  user = request.session["user"]
except KeyError:
  # redirect user to login page

# check for admin
if not user.isAdmin():
  # display error page

# show admin page
...

Here is a version using decorators

# login decorator checks whether the user is logged in
def login_required(fn):
  def _check(request, *args, **kwargs):
    try:
      user = request.session["user"]
    except KeyError:
      # redirect to login page
    # user is logged in, call the function
    return fn(args, kwargs)
  return _check

# admin decorator checks whether the user is an admin
def admin_required(fn):
  def _check(request, *args, **kwargs):
    user = request.session["user"]
    if not user.isAdmin():
      # return the error page
    # user is admin, call the function
    return fn(args, kwargs)
  return _check

@login_required
@admin_required
def admin(request):
  # show admin page
  ...

This is how it works. We first have the admin function. All it does is implement the admin code. We decorate it with the admin_required and login_required decorators. When the admin function is called, it first enters the login_required decorator function which checks for login. If the user is not logged in, it redirects to the login page, else it calls the function. The function passed to login_required is the admin_required decorated function. So if login passes, it calls the admin_required decorator, checks for admin. If the user is not an admin, it displays an error message, else calls the function, which in this case is the original admin function.

See how neatly the guards are separated from the core logic. The admin function contains only the core logic, while the list of guards is neatly arranged as decorators. This method also has another advantage — it is easy to apply the same guards to other functions. Take a look at this

# no decorators
def login(request):
  ...

@login_required
@admin_required
def admin(request):
  ...

@login_required
def dashboard(request):
  ...

@login_required
@project_permission_required
def view_project(request, project):
  ...

Now each function only implements the core logic, while all the guard logic is taken care of by the decorators. It is easy to see the core logic and easy to see the guard conditions applied for the function. If I want some other function to have a guard, I can just add a decorator to it without touching the core logic. Best of all, the code is self descriptive and very easy to read.

This post is a part of the selected archive.