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

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 2

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()