timer.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. import time
  2. import argparse
  3. class TimerSubcategory:
  4. def __init__(self, timer, category):
  5. self.timer = timer
  6. self.category = category
  7. self.start = None
  8. self.original_base_category = timer.base_category
  9. def __enter__(self):
  10. self.start = time.time()
  11. self.timer.base_category = self.original_base_category + self.category + "/"
  12. self.timer.subcategory_level += 1
  13. if self.timer.print_log:
  14. print(f"{' ' * self.timer.subcategory_level}{self.category}:")
  15. def __exit__(self, exc_type, exc_val, exc_tb):
  16. elapsed_for_subcategroy = time.time() - self.start
  17. self.timer.base_category = self.original_base_category
  18. self.timer.add_time_to_record(self.original_base_category + self.category, elapsed_for_subcategroy)
  19. self.timer.subcategory_level -= 1
  20. self.timer.record(self.category, disable_log=True)
  21. class Timer:
  22. def __init__(self, print_log=False):
  23. self.start = time.time()
  24. self.records = {}
  25. self.total = 0
  26. self.base_category = ''
  27. self.print_log = print_log
  28. self.subcategory_level = 0
  29. def elapsed(self):
  30. end = time.time()
  31. res = end - self.start
  32. self.start = end
  33. return res
  34. def add_time_to_record(self, category, amount):
  35. if category not in self.records:
  36. self.records[category] = 0
  37. self.records[category] += amount
  38. def record(self, category, extra_time=0, disable_log=False):
  39. e = self.elapsed()
  40. self.add_time_to_record(self.base_category + category, e + extra_time)
  41. self.total += e + extra_time
  42. if self.print_log and not disable_log:
  43. print(f"{' ' * self.subcategory_level}{category}: done in {e + extra_time:.3f}s")
  44. def subcategory(self, name):
  45. self.elapsed()
  46. subcat = TimerSubcategory(self, name)
  47. return subcat
  48. def summary(self):
  49. res = f"{self.total:.1f}s"
  50. additions = [(category, time_taken) for category, time_taken in self.records.items() if time_taken >= 0.1 and '/' not in category]
  51. if not additions:
  52. return res
  53. res += " ("
  54. res += ", ".join([f"{category}: {time_taken:.1f}s" for category, time_taken in additions])
  55. res += ")"
  56. return res
  57. def dump(self):
  58. return {'total': self.total, 'records': self.records}
  59. def reset(self):
  60. self.__init__()
  61. parser = argparse.ArgumentParser(add_help=False)
  62. parser.add_argument("--log-startup", action='store_true', help="print a detailed log of what's happening at startup")
  63. args = parser.parse_known_args()[0]
  64. startup_timer = Timer(print_log=args.log_startup)
  65. startup_record = None