signal_timer context manager

signal_timer.py

import signal
import time

class TimerException(Exception): pass

from contextlib import contextmanager

@contextmanager
def signal_timer(mili):
    old = signal.getitimer(signal.ITIMER_REAL)
    start = time.time()
    tick = float(mili) / 1000.
    def timeout(signum, frame):
        code = frame.f_code
        duration = (time.time() - start) * 1000
        signal.setitimer(signal.ITIMER_REAL, 0, 0)  # don't race to the reset, we might lose
        raise TimerException("Function ran too long %s:%s <%s> (%dms)" % (code.co_name,
                                                                          frame.f_lineno,
                                                                          code.co_filename,
                                                                          duration))
    signal.signal(signal.SIGALRM, timeout)
    signal.setitimer(signal.ITIMER_REAL, tick, tick)
    try:
        yield
    except TimerException, e:
        print str(e)
    duration = (time.time() - start)
    signal.setitimer(signal.ITIMER_REAL, max(0, old[0] - duration), old[1])
    

with signal_timer(500):
    with signal_timer(50):
               i = 0
               while True:
                    i += 1
                    if i % 100000 == 0:
                        print "yes"
    while True:
        i += 1

yes
yes
Function ran too long module:37 testsignal.py (50ms)
Function ran too long module:39 testsignal.py (500ms)