Using decorator in a class to access an attribute in self

Question

I have a class which provides some generic structure for a sequential data processing pipeline. I would like to time the execution of a method, and save that to a dictionary attribute of self (self.timings).

from functools import wraps
import time

class Pipeline(object):

    def __init__(self):
        self.steps = {}
        self.timings = {}

    # Decorator for adding functions to pipeline
    def step(self, step_name):
        def step_decorator(f):
            self.steps[step_name] = f
        return step_decorator

    # Decorator for timing a step
    def time_step(f):
        @wraps(f)
        def timed(*args, **kwargs):
            start = time.time()
            result = f(*args, **kwargs)
            end = time.time()
            self.timings[f.__name__] = end - start
            return result
        return timed

    @time_step
    def example_method(self):
        if 'example_func' in self.steps:
            self.output = self.steps['example_func']()

I can instantiate a Pipeline and add a step to it:

pipeline = Pipeline()

@pipeline.step('example_func')
def example_func():
    for i in range(10000):
        pass
    return 'Completed!'

But when I attempt to run pipeline.example_method(), it cannot access self:

pipeline.example_method()
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-54-8204774a5649> in <module>()
----> 1 pipeline.example_method()

<ipython-input-51-ffb2e95a110a> in timed(*args, **kwargs)
     21             result = f(*args, **kwargs)
     22             end = time.time()
---> 23             self.timings[f.__name__] = end - start
     24             return result
     25         return timed

NameError: name 'self' is not defined

I have tried adding self to the parameters in the time_step definition, but that causes another error. Is there a straightforward way of accessing attributes from a decorated method?

Answer 1

Your @time_step() decorator is not a bound method, when @time_step runs, it is just a function object. No self is defined in the function, nor is it defined in the wrapper that this decorator returns.

If you use time_step() only on methods, then you can count on the returned wrapper function being bound (it's just another function object in the class body in that case, so is treated like any other function attached to a class and looked up on the instance):

def time_step(f):
    @wraps(f)
    def timed(self, *args, **kwargs):
        start = time.time()
        result = f(self, *args, **kwargs)
        end = time.time()
        self.timings[f.__name__] = end - start
        return result
    return timed

Note that you do have to pass the self argument on to the f() call, since f is also not bound.