Decorator to time specific lines of the code instead of whole method?

Question

Lets assume a simple method :

def test_method():
    a = 1
    b = 10000
    c = 20000
    sum1 = sum(range(a,b))
    sum2 = sum(range(b,c))
    return (sum1,sum2)

To time this method using a decorator, a simple decorator would be :

from functools import wraps
def timed_decorator(f):
    @wraps(f)
    def wrapper(*args, **kwds):
        start = time.time()
        result = f(*args, **kwds)
        elapsed = (time.time() - start)*1000
        logger.debug("f::{0} t::{1:0.2f} ms".format(f.__name__, elapsed))
        return result
    return wrapper

Now if I want to time specific lines of test_method say line 4 sum1 = sum(range(a,b)) , the current implementation involves inline coding like:

 def test_method():
        a = 1
        b = 10000
        c = 20000
        start = time.time()
        sum1 = sum(range(a,b)) # timing specific line or lines
        elapsed = (time.time() - start)*1000
        logger.debug("This part took::{1:0.2f} ms".format(elapsed))
        sum2 = sum(range(b,c))
        return (sum1,sum2)

The intention is to use the decorator to time lines M to N of a specific method without modifying the code in the method. Is it possible to inject such logic using a decorator ?

Answer 1

Very simple solution with a custom context manager:

class elapsed:
    def __enter__(self): self.start = time.time()
    def __exit__(self, *args): print("%.1f ms" % ((time.time() - self.start)*1000))

Example usage:

with elapsed():
    sum1 = sum(x ** 2 for x in range(1, 1000000))
# 547.0 ms

More about this: Decorator-like syntax for a specific line of code

---

Another solution: here is a slight variation of @NiklasR's answer without logger but print, and a ready-to-run example:

import contextlib, time

@contextlib.contextmanager
def time_measure(ident):
    tstart = time.time()
    yield
    elapsed = time.time() - tstart
    print("{0}: {1} ms".format(ident, elapsed))

with time_measure('hello'):
    sum1 = sum(x ** 2 for x in range(1, 1000000))

# hello: 0.577033281326294 ms

Answer 2

It's pretty ugly, and not very stable code. but the only way I found to do this task is to exec the code of the function again, after injecting your code.
Something like this:

import inspect
import re
import time

def inject_timer(f,n,m):
    codelines = inspect.getsourcelines(f)[0]
    ident_lvl = re.search("^[ \t]*",codelines[n]).group(0)
    codelines.insert(n,ident_lvl + "start_longJibrishTo_preventCollision = time.time()\n")
    codelines.insert(m+2,ident_lvl + "elapsed_longJibrishTo_preventCollision = (time.time() - start_longJibrishTo_preventCollision)*1000\n")
    codelines.insert(m+3,ident_lvl + """print("f::{0} t::{1:0.2f} ms".format("""+f.__name__+""", elapsed_longJibrishTo_preventCollision))\n""")
    #print "".join(codelines)
    exec "".join(codelines) in globals()

def test_method():
    a = 1
    b = 10000
    time.sleep(2)
    c = 20000    
    sum1 = sum(range(a,b))
    sum2 = sum(range(b,c))    
    return (sum1,sum2)

inject_timer(test_method,3,5)

Answer 3

One way I can think of is to use sys.settrace() and record time when handling "line" event in the tracer function. But one caveat is, the practice of setting a tracer may cause the time recorded to be inaccurate.

The general idea is:

1. Set a tracer function in the decorator that wraps the target method.
2. Get the line number for the first line of this method, with FLN = inspect.currentframe().f_lineno.
3. In the tracer function, handle "call" event and return a local tracer function to trace the "line" events in the scope. Read this if you are confused.
4. Within the local tracer function, get the current line number LN, if LN-FLN == M, record the start time; if LN-FLN == N, record the end time, the time taken to execute lines M to N is endtime - starttime.

code:

import sys
from functools import wraps
import time
import linecache

_func_name_ = None
_func_ln_ = 0
_start_ = 0
_end_ = 0
_timestamp_ = 0

def trace_calls(frame, event, arg): 
    global _func_name_, _func_ln_
    def trace_lines(frame, event, arg): 
        global _timestamp_
            if event != 'line':
                return
        line_no = frame.f_lineno
        filename = frame.f_code.co_filename
        if line_no-_func_ln_ == _start_:                        
            _timestamp_ = time.time()
            print "%d %s TS:%d"%(line_no, linecache.getline(filename, line_no)[:-1], _timestamp_)

        elif line_no-_func_ln_ == _end_:
            _timestamp_ = time.time() - _timestamp_
            print "%d %s"%(line_no, linecache.getline(filename, line_no)[:-1])
            print "Lines %d to %d of %s takes %d seconds."%(_start_, _end_, _func_name_,  _timestamp_)      

    if event != 'call':    
        return   

    co = frame.f_code      
    _func_ln_ = frame.f_lineno  # record the line number at function entry point
    func_name = co.co_name 

    if func_name != _func_name_:
        return             
    return trace_lines

def time_lines(start, end):
    global _start_, _end_
    _start_, _end_ = start+1, end+2     # function name takes a line, end is inclusive
    def inner(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            global _func_name_
            _func_name_ = f.__name__    
            sys.settrace(trace_calls)
            f(*args, **kwargs)
            sys.settrace(None)
        return wrapper
    return inner

@time_lines(2,4)
def tested_func():
    print "Enter target function"
    time.sleep(2)
    time.sleep(1)
    time.sleep(3)
    print "Exit target function"

if __name__=="__main__":
    tested_func()

Answer 4

You can use a context manager.

import contextlib

@contextlib.contextmanager
def time_measure(ident):
    tstart = time.time()
    yield
    elapsed = time.time() - tstart
    logger.debug("{0}: {1} ms".format(ident, elapsed))

In your code, you use it like

with time_measure('test_method:sum1'):
    sum1 = sum(range(a, b))

---

By the way, if you want to improve your code, you can use the Gaussian Sum Formula (explained here) instead of sum(range(a, b)).

def sum_range(a, b):
    r_a = (a ** 2 + a) / 2 - a
    r_b = (b ** 2 + b) / 2 - b
    return r_b - r_a

Answer 5

A decorator can only decorate callables (e.g. functions, methods, classes). A single line or a group of lines are not callable as long as you do not wrap them in their own callable.

For timing a unit of your code you should choose an appropriate number of repetitions. The goal is to make sure that the execution time is longer than just a few micro or milliseconds, otherwise the measurement error will be too large.

Did you have a look at the timeit module?