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