import unittest
import atexit
import time
import math
import sched
from collections import defaultdict
from contextlib import contextmanager

_timethis_stats = defaultdict(list)
_flag = False

def timethis(what):
    global _flag
    def _printstats():
        if not _timethis_stats:
            return
        maxwidth = max(len(str(key)) for key in _timethis_stats)
        items = sorted(_timethis_stats.items(), key=lambda x: str(x[0]))
        for key, times in items:
            name = getattr(key, '__name__', str(key))
            # compute average and standard deviation
            mean = sum(times) / float(len(times))
            stddev = math.sqrt(sum((x-mean)**2 for x in times) / len(times))
            print("%-30s : time=%0.5f : calls=%s : stdev=%0.5f" % \
                    (name, mean, len(times), stddev))

    if not _flag:
        _flag = True
        atexit.register(_printstats)
    @contextmanager
    def benchmark():
        start = time.time()
        yield
        end = time.time()
        _timethis_stats[what].append(end-start)


    if hasattr(what,"__call__"):
        def timed(*args,**kwargs):
            with benchmark():
                return what(*args,**kwargs)
        return timed
    else:
        return benchmark()


class Benchmark(unittest.TestCase):
    times = 1000
    noop = lambda: 0

    def setUp(self):
        self.scheduler = sched.scheduler()
        self.events = []
        self.populate()

    def tearDown(self):
        del self.events[:]

    def populate(self):
        for x in range(self.times):
            ev = self.scheduler.enter(0, 0, lambda: 0)
            self.events.append(ev)

    @timethis
    def test_run(self):
        self.scheduler.run()

    @timethis
    def test_enter(self):
        for x in range(self.times):
            self.scheduler.enter(0, 0, lambda: 0)

    @timethis
    def test_cancel(self):
        for ev in self.events:
            self.scheduler.cancel(ev)

    @timethis
    def test_queue(self):
        for x in range(self.times):
            list(self.scheduler.queue)

    @timethis
    def test_empty(self):
        for x in range(self.times):
            self.scheduler.empty()

unittest.main()