]>
Commit | Line | Data |
---|---|---|
53e6db90 DC |
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) |