debug.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. import os
  2. import time
  3. from contextlib import contextmanager
  4. from typing import Callable, Optional
  5. _inited = False
  6. def _lazy_colorama_init():
  7. """
  8. Lazily init colorama if necessary, not to screw up stdout if debugging is
  9. not enabled.
  10. This version of the function does nothing.
  11. """
  12. try:
  13. if os.name == 'nt':
  14. # Does not work on Windows, as pyreadline and colorama interfere
  15. raise ImportError
  16. else:
  17. # Use colorama for nicer console output.
  18. from colorama import Fore, init # type: ignore[import]
  19. from colorama import initialise
  20. def _lazy_colorama_init(): # noqa: F811
  21. """
  22. Lazily init colorama if necessary, not to screw up stdout is
  23. debug not enabled.
  24. This version of the function does init colorama.
  25. """
  26. global _inited
  27. if not _inited:
  28. # pytest resets the stream at the end - causes troubles. Since
  29. # after every output the stream is reset automatically we don't
  30. # need this.
  31. initialise.atexit_done = True
  32. try:
  33. init(strip=False)
  34. except Exception:
  35. # Colorama fails with initializing under vim and is buggy in
  36. # version 0.3.6.
  37. pass
  38. _inited = True
  39. except ImportError:
  40. class Fore: # type: ignore[no-redef]
  41. RED = ''
  42. GREEN = ''
  43. YELLOW = ''
  44. MAGENTA = ''
  45. RESET = ''
  46. BLUE = ''
  47. NOTICE = object()
  48. WARNING = object()
  49. SPEED = object()
  50. enable_speed = False
  51. enable_warning = False
  52. enable_notice = False
  53. # callback, interface: level, str
  54. debug_function: Optional[Callable[[str, str], None]] = None
  55. _debug_indent = 0
  56. _start_time = time.time()
  57. def reset_time():
  58. global _start_time, _debug_indent
  59. _start_time = time.time()
  60. _debug_indent = 0
  61. def increase_indent(func):
  62. """Decorator for makin """
  63. def wrapper(*args, **kwargs):
  64. with increase_indent_cm():
  65. return func(*args, **kwargs)
  66. return wrapper
  67. @contextmanager
  68. def increase_indent_cm(title=None, color='MAGENTA'):
  69. global _debug_indent
  70. if title:
  71. dbg('Start: ' + title, color=color)
  72. _debug_indent += 1
  73. try:
  74. yield
  75. finally:
  76. _debug_indent -= 1
  77. if title:
  78. dbg('End: ' + title, color=color)
  79. def dbg(message, *args, color='GREEN'):
  80. """ Looks at the stack, to see if a debug message should be printed. """
  81. assert color
  82. if debug_function and enable_notice:
  83. i = ' ' * _debug_indent
  84. _lazy_colorama_init()
  85. debug_function(color, i + 'dbg: ' + message % tuple(repr(a) for a in args))
  86. def warning(message, *args, format=True):
  87. if debug_function and enable_warning:
  88. i = ' ' * _debug_indent
  89. if format:
  90. message = message % tuple(repr(a) for a in args)
  91. debug_function('RED', i + 'warning: ' + message)
  92. def speed(name):
  93. if debug_function and enable_speed:
  94. now = time.time()
  95. i = ' ' * _debug_indent
  96. debug_function('YELLOW', i + 'speed: ' + '%s %s' % (name, now - _start_time))
  97. def print_to_stdout(color, str_out):
  98. """
  99. The default debug function that prints to standard out.
  100. :param str color: A string that is an attribute of ``colorama.Fore``.
  101. """
  102. col = getattr(Fore, color)
  103. _lazy_colorama_init()
  104. print(col + str_out + Fore.RESET)