Question

I’m trying to design an interface that tests whether a user is logged in before running certain functions in the class. Rather than:

class UserDoesStuff(object):
    def doIfLoggedIn(self):
        if self.checkLogin():
           [...do the stuff...]

I was wondering if I could have something like this:

def protected(self):
     if not self.checkLogin():
         raise UserLoginError()

@protected
def doIfLoggedIn(self):
    [...do the stuff...]

This of course doesn’t work, but is there a way to do this using decorators?

Answer 1

Decorators (the simplest ones, with no extra args) expect functions as input:

import functools

def protected(fun):
    @functools.wraps(fun)
    def wrapper(self, *args, **kwargs):
        if not self.checkLogin():
            raise UserLoginError()
        return fun(self, *args, **kwargs)

    return wrapper # this is what replaces the original method

@protected
def doIfLoggedIn(self):
    ...