| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132 |
- import os
- import time
- from contextlib import contextmanager
- from typing import Callable, Optional
- _inited = False
- def _lazy_colorama_init():
- """
- Lazily init colorama if necessary, not to screw up stdout if debugging is
- not enabled.
- This version of the function does nothing.
- """
- try:
- if os.name == 'nt':
- # Does not work on Windows, as pyreadline and colorama interfere
- raise ImportError
- else:
- # Use colorama for nicer console output.
- from colorama import Fore, init # type: ignore[import]
- from colorama import initialise
- def _lazy_colorama_init(): # noqa: F811
- """
- Lazily init colorama if necessary, not to screw up stdout is
- debug not enabled.
- This version of the function does init colorama.
- """
- global _inited
- if not _inited:
- # pytest resets the stream at the end - causes troubles. Since
- # after every output the stream is reset automatically we don't
- # need this.
- initialise.atexit_done = True
- try:
- init(strip=False)
- except Exception:
- # Colorama fails with initializing under vim and is buggy in
- # version 0.3.6.
- pass
- _inited = True
- except ImportError:
- class Fore: # type: ignore[no-redef]
- RED = ''
- GREEN = ''
- YELLOW = ''
- MAGENTA = ''
- RESET = ''
- BLUE = ''
- NOTICE = object()
- WARNING = object()
- SPEED = object()
- enable_speed = False
- enable_warning = False
- enable_notice = False
- # callback, interface: level, str
- debug_function: Optional[Callable[[str, str], None]] = None
- _debug_indent = 0
- _start_time = time.time()
- def reset_time():
- global _start_time, _debug_indent
- _start_time = time.time()
- _debug_indent = 0
- def increase_indent(func):
- """Decorator for makin """
- def wrapper(*args, **kwargs):
- with increase_indent_cm():
- return func(*args, **kwargs)
- return wrapper
- @contextmanager
- def increase_indent_cm(title=None, color='MAGENTA'):
- global _debug_indent
- if title:
- dbg('Start: ' + title, color=color)
- _debug_indent += 1
- try:
- yield
- finally:
- _debug_indent -= 1
- if title:
- dbg('End: ' + title, color=color)
- def dbg(message, *args, color='GREEN'):
- """ Looks at the stack, to see if a debug message should be printed. """
- assert color
- if debug_function and enable_notice:
- i = ' ' * _debug_indent
- _lazy_colorama_init()
- debug_function(color, i + 'dbg: ' + message % tuple(repr(a) for a in args))
- def warning(message, *args, format=True):
- if debug_function and enable_warning:
- i = ' ' * _debug_indent
- if format:
- message = message % tuple(repr(a) for a in args)
- debug_function('RED', i + 'warning: ' + message)
- def speed(name):
- if debug_function and enable_speed:
- now = time.time()
- i = ' ' * _debug_indent
- debug_function('YELLOW', i + 'speed: ' + '%s %s' % (name, now - _start_time))
- def print_to_stdout(color, str_out):
- """
- The default debug function that prints to standard out.
- :param str color: A string that is an attribute of ``colorama.Fore``.
- """
- col = getattr(Fore, color)
- _lazy_colorama_init()
- print(col + str_out + Fore.RESET)
|