A method timing decorator for Python

Here is a simple method decorator for tracking the time spent in each method of an instance, that turned out to be useful for debugging some slow code:

def timer_wrap(func):
    def wrapper(*args, **kwargs):
        self = args[0]
        start_t = time.time()
        ans = func(*args, **kwargs)
        stop_t = time.time()
        name = func.__name__
        if not hasattr(self, "_timers"):
            setattr(self, "_timers", {})
        if not name in self._timers:
            self._timers[name] = 0
        self._timers[name] += stop_t - start_t
        return ans
    return wrapper

The decorator works by storing the timing results of each method in a _timers dictionary in the object (which is created if it doesn’t exist). This dictionary is a mapping from the method name to the total time spent in that method. One caveat is that if one method calls another method, the time will be added to both (for my application this wasn’t an issue, but your mileage may vary).

Below is a simple example to show how this could be used in practice:

import time

def timer_wrap(func):
    # as above

class A(object):
    def __init__(self):
        pass

    @timer_wrap
    def one(self):
        time.sleep(1)

    @timer_wrap
    def two(self):
        time.sleep(2)

def main():
    a = A()
    a.one()
    a.two()
    print(a._timers)

if __name__ == '__main__':
    main()

this gives:

{'one': 1.0010690689086914, 'two': 2.0021371841430664}

Of course, after I wrote this I found out that this had been done before, but oh well.

Update: If you do care about methods that call each other and want to keep track how much time is spent in each separately, you can use the following code:

class Timer(object):
    def __init__(self):
        self.timers = {}
        self._stack = []
        self.start = None

    def add_to_timer(self, name, duration):
        if not name in self.timers:
            self.timers[name] = 0
        self.timers[name] += duration

    def stack(self, name):
        # stop running timer, start new timer, add name to stack
        if len(self._stack):
            self.add_to_timer(self._stack[0], time.time() - self.start)
        self.start = time.time()
        self._stack.insert(0, name)

    def pop(self):
        # pop name from stack, restart previous timer
        self.add_to_timer(self._stack.pop(0), time.time() - self.start)
        self.start = time.time()

TIMER = Timer()

def timer_wrap(func):
    def wrapper(*args, **kwargs):
        name = func.__name__
        TIMER.stack(name)
        ans = func(*args, **kwargs)
        TIMER.pop()
        return ans
    return wrapper

This works by keeping a stack of timers that are stopped and started at each method that is decorated with timer_wrap. This allows you to measure the time spent in each method separately, provided that you use the decorator when you want to get fine-grained control. Note that it isn’t very clean to use a global timer object, but for debugging purposes this works fine.