Search code examples
pythonpython-decorators

How to pass an instance variable to a decorator inside class definition?


I'd like to use a decorator which accepts an argument, checks if that argument is not None, and if True it lets the decorated function run.

I want to use this decorator inside a class definition, because I have a set of class methods which starts with checking if a specific instance variable is None or not. I think it would look nicer if I used a decorator.

I'd like to do something like this:

# decorator
def variable_tester(arg):
    def wrap(f):
        def wrapped_f(*args):
            if arg is not None:
                f(*args)
            else:
                pass
        return wrapped_f
    return wrap

# class definition  
class my_class(object):
    def __init__(self):
        self.var = None

    @variable_tester(self.var) # This is wrong. How to pass self.var to decorator?
    def printout(self):
        print self.var

    def setvar(self, v):
        self.var = v

# testing code
my_instance = my_class()
my_instance.printout() # It should print nothing

my_instance.setvar('foobar')
my_instance.printout() # It should print foobar

Solution

  • This is a little tricky, because you want to do several things, each of them a little finicky: (1) you want to pass an argument to a decorator and (2) you want that argument to refer to the instance, but the instance doesn't exist at the time of decoration, so we'll have to defer it somehow. You could use a function, or an itemgetter, but here I'll use a string, and I won't use functools.wraps like I should 'cause I'm lazy.

    Something like:

    # really, it's variable-tester-factory
    def variable_tester(target):
        def deco(function):
            def inner(self, *args, **kwargs):
                if getattr(self, target) is not None:
                    return function(self, *args, **kwargs)
            return inner
        return deco
    
    class my_class(object):
        def __init__(self):
            self.var = None
    
        @variable_tester('var')
        def printout(self):
            print self.var
    

    should work:

    >>> a = my_class()
    >>> print a.var
    None
    >>> a.printout()
    >>> a.var = 'apple'
    >>> a.printout()
    apple