indico/util/benchmark.py

   1 import time
   2 from math import isinf
   3
   4 from indico.util.console import colored
   5
   6
   7 class Benchmark(object):
   8     """Simple benchmark class
   9
  10     Can be used manually or as a contextmanager:
  11
  12     with Benchmark() as b:
  13         do_stuff()
  14     b.print_result()
  15     """
  16     def __init__(self, start=False):
  17         self._start_time = None
  18         self._end_time = None
  19         if start:
  20             self.start()
  21
  22     def start(self):
  23         self._start_time = time.time()
  24         return self
  25
  26     def stop(self):
  27         self._end_time = time.time()
  28
  29     def __enter__(self):
  30         return self.start()
  31
  32     def __exit__(self, exc_type, exc_val, exc_tb):
  33         self.stop()
  34
  35     def __float__(self):
  36         if self._start_time is None:
  37             return float('-inf')  # not started
  38         elif self._end_time is None:
  39             return float('inf')  # not finished
  40         return self._end_time - self._start_time
  41
  42     def __str__(self):
  43         duration = float(self)
  44         if isinf(duration):
  45             return str(duration)
  46         return '{:.05f}'.format(duration)
  47
  48     __repr__ = __str__
  49
  50     def print_result(self, slow=float('inf'), veryslow=float('inf')):
  51         duration = float(self)
  52         if duration == float('-inf'):
  53             print colored('skipped', 'blue', attrs=['bold'])
  54         elif duration == float('inf'):
  55             print colored('running', 'red')
  56         elif duration >= veryslow:
  57             print colored(str(self), 'red', attrs=['bold'])
  58         elif duration >= slow:
  59             print colored(str(self), 'yellow', attrs=['bold'])
  60         else:
  61             print colored(str(self), 'green', attrs=['bold'])