Poor man's profiler

pmp.py

# Poor man's profiler
# (c) 2013 Jasper Spaans, covered by the WTFPL [wtfpl.org]

import collections
import signal
import threading
import time

local_context_stack = threading.local()

class Profiler(object):
    """Context manager for profiling your data"""
    def __init__(self, interval=0.01):
        self.interval = interval
        self.profiler_data = threading.local()

    def __enter__(self):
        delay, interval = signal.getitimer(signal.ITIMER_VIRTUAL)
        if delay and interval:
            raise RuntimeError("Refusing to run while another user of itimer is active")
        signal.signal(signal.SIGVTALRM, self.handle_vtalrm)
        signal.setitimer(signal.ITIMER_VIRTUAL, self.interval, self.interval)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        signal.setitimer(signal.ITIMER_VIRTUAL, 0, 0)

    @property
    def _count(self):
        try:
            count = self.profiler_data.count
        except AttributeError:
            count = self.profiler_data.count = collections.Counter()
        return count

    def get_data(self):
        return self._count

    def handle_vtalrm(self, signum, stackframe):
        try:
            context_name = local_context_stack.data[-1]
        except AttributeError:
            context_name = "None"
        self._count[context_name] += 1


class ProfilerContext(object):
    """Profiler context - use this to declare the context of the currently running function"""
    def __init__(self, context_name):
        self.context_name = context_name

    def __enter__(self):
        try:
            stack = local_context_stack.data
        except AttributeError:
            stack = local_context_stack.data = []
        stack.append(self.context_name)

    def __exit__(self, exc_type, exc_val, exc_tb):
        local_context_stack.data.pop(-1)


with Profiler() as p:
    with ProfilerContext("c1"):
        s = ""
        for t in range(100000):
            with ProfilerContext("c2"):
                s = s + "a"
            s = s + "b"
        print p.get_data()