userver: /home/antonyzhilin/arcadia/taxi/uservices/userver/testsuite/pytest_plugins/pytest_userver/utils/colorize.py Source File
All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Modules Pages Concepts
colorize.py
1import argparse
2import contextlib
3import enum
4import itertools
5import json
6import sys
7
8from . import tskv
9
10
11class ColorArg(enum.Enum):
12 AUTO = 'auto'
13 ALWAYS = 'always'
14 NEVER = 'never'
15
16
17# Info about colors https://misc.flogisoft.com/bash/tip_colors_and_formatting
18class Colors:
19 BLACK = '\033[30m'
20 RED = '\033[31m'
21 GREEN = '\033[32m'
22 YELLOW = '\033[33m'
23 BLUE = '\033[34m'
24 GRAY = '\033[37m'
25 DARK_GRAY = '\033[37m'
26 BRIGHT_RED = '\033[91m'
27 BRIGHT_GREEN = '\033[92m'
28 BRIGHT_YELLOW = '\033[93m'
29 DEFAULT = '\033[0m'
30 DEFAULT_BG = '\033[49m'
31 BG_BLACK = '\033[40m'
32
33 # No bright red color, no close colors, no too dark colors
34 __NICE_COLORS = [
35 '\033[38;5;{}m'.format(x)
36 for x in itertools.chain(
37 range(2, 7),
38 range(10, 15),
39 range(38, 51, 2),
40 range(75, 87, 2),
41 range(112, 123, 3),
42 range(128, 159, 3),
43 range(164, 195, 3),
44 range(203, 231, 3),
45 )
46 ]
47
48 @staticmethod
49 def colorize(value):
50 ind = hash(value) % len(Colors.__NICE_COLORS)
51 return Colors.__NICE_COLORS[ind]
52
53
54LEVEL_COLORS = {
55 'TRACE': Colors.DARK_GRAY,
56 'DEBUG': Colors.GRAY,
57 'INFO': Colors.GREEN,
58 'WARNING': Colors.YELLOW,
59 'ERROR': Colors.RED,
60 'CRITICAL': Colors.BRIGHT_RED,
61 'none': Colors.DEFAULT,
62}
63HTTP_STATUS_COLORS = {
64 '2': Colors.GREEN,
65 '3': Colors.GREEN,
66 '4': Colors.YELLOW,
67 '5': Colors.RED,
68}
69
70HTTP_LOCALHOST_PREFIX = 'http://localhost'
71
72
74 def __init__(self, *, verbose=False, colors_enabled=True):
75 self._requests = {}
76 self.verbose = verbose
77 self.colors_enabled = colors_enabled
78
79 def colorize_line(self, line):
80 if not line.startswith('tskv\t'):
81 return line
82 return self.colorize_tskv(line)
83
84 def colorize_tskv(self, line):
85 row = tskv.parse_line(line)
86 return self.colorize_row(row)
87
88 def colorize_row(self, row):
89 row = row.copy()
90 flowid = '-'.join([row.get(key, '') for key in ('link', 'trace_id')])
91
92 entry_type = row.pop('_type', None)
93 link = row.get('link', None)
94 level = row.pop('level', 'none')
95 text = row.pop('text', '')
96
97 extra_fields = []
98 if entry_type == 'request':
99 self._requests[link] = self._build_request_info(row)
100 if 'body' in row:
101 extra_fields.append(
102 'request_body='
103 + self.textcolor(
104 try_reformat_json(row.pop('body')),
105 Colors.YELLOW,
106 ),
107 )
108 elif entry_type == 'response':
109 if 'body' not in row:
110 raise RuntimeError(
111 f'Response log record without "body" tag. Looks like in the C++ code the tracing::Span of a'
112 f'request was moved out or corrupted. Link: {link}. Text: {text}. Other: {row}'
113 )
114 if 'meta_code' in row:
115 status_code = row.pop('meta_code')
116 extra_fields.append(
117 self._http_status('meta_code', status_code),
118 )
119 if not text:
120 text = 'Response finished'
121 extra_fields.append(
122 'response_body='
123 + self.textcolor(
124 try_reformat_json(row.pop('body')),
125 Colors.YELLOW,
126 ),
127 )
128 elif entry_type == 'mockserver_request':
129 text = 'Mockserver request finished'
130 if 'meta_code' in row:
131 status_code = str(row.pop('meta_code'))
132 extra_fields.append(
133 self._http_status('meta_code', status_code),
134 )
135 for key in ('method', 'url', 'status', 'exc_info', 'delay'):
136 value = row.pop(key, None)
137 if value:
138 extra_fields.append(f'{key}={value}')
139
140 if link in self._requests:
141 logid = f'[{self._requests[link]}]'
142 elif link is not None:
143 logid = f'[{link}]'
144 else:
145 logid = '<userver>'
146
147 level_color = LEVEL_COLORS.get(level)
148 flow_color = Colors.colorize(flowid)
149
150 fields = [
151 self.textcolor(f'{level:<8}', level_color),
152 self.textcolor(logid, flow_color),
153 ]
154 if text:
155 http_url_info = row.pop('http_url', None)
156 if http_url_info:
157 http_url_info = http_url_info.removeprefix(
158 HTTP_LOCALHOST_PREFIX,
159 )
160 http_url_info = http_url_info.removesuffix('?')
161 http_url_info = http_url_info[http_url_info.find('/') + 1 :]
162
163 meta_code = row.pop('meta_code', None)
164 if meta_code:
165 http_url_info += f' meta_code={meta_code}'
166
167 text = f'{self.textcolor(http_url_info, Colors.GREEN)} {text}'
168
169 fields.append(text)
170 elif self.verbose:
171 fields.append('<NO TEXT>')
172 else:
173 return None
174
175 fields.extend(extra_fields)
176 if self.verbose:
177 fields.extend([f'{k}={v}' for k, v in row.items()])
178 return ' '.join(fields)
179
180 def textcolor(self, text, color):
181 if not self.colors_enabled:
182 return str(text)
183 return f'{color}{text}{Colors.DEFAULT}'
184
185 def _http_status(self, key, status):
186 color = HTTP_STATUS_COLORS.get(status[:1], Colors.DEFAULT)
187 return self.textcolor(f'{key}={status}', color)
188
189 def _build_request_info(self, row):
190 if 'uri' not in row:
191 return None
192 uri = row['uri']
193 method = row.get('method', 'UNKNOWN')
194 return f'{method} {uri}'
195
196
197def format_json(obj):
198 encoded = json.dumps(
199 obj,
200 indent=2,
201 separators=(',', ': '),
202 sort_keys=True,
203 ensure_ascii=False,
204 )
205 return encoded
206
207
208def try_reformat_json(body):
209 try:
210 # TODO: unescape string
211 data = json.loads(body)
212 return format_json(data)
213 except ValueError:
214 return body
215
216
217def colorize(stream, verbose=False, colors_enabled=True):
218 colorizer = Colorizer(verbose=verbose, colors_enabled=colors_enabled)
219 for line in stream:
220 line = line.rstrip('\r\n')
221 color_line = colorizer.colorize_line(line)
222 if color_line is not None:
223 print(color_line)
224
225
226def parse_color(value):
227 if value in ('always', 'force', 'yes', 'enable'):
228 return ColorArg.ALWAYS
229 if value in ('never', 'no', 'disable'):
230 return ColorArg.NEVER
231 if value == 'auto':
232 return ColorArg.AUTO
233 raise ValueError(f'Unknown color option {value!r}')
234
235
236def colorize_main():
237 parser = argparse.ArgumentParser(description='Colorize userver log file.')
238 parser.add_argument(
239 '--verbose',
240 '-v',
241 action='store_true',
242 help='Be verbose',
243 )
244 parser.add_argument(
245 '--color',
246 metavar='WHEN',
247 help=('Control color highlighting, WHEN is always, never or auto (default)'),
248 nargs='?',
249 type=parse_color,
250 default=ColorArg.AUTO,
251 const=ColorArg.ALWAYS,
252 )
253 parser.add_argument(
254 '--no-color',
255 help='Turn off color highlighting',
256 dest='color',
257 action='store_const',
258 const=ColorArg.NEVER,
259 )
260 parser.add_argument(
261 'log',
262 help='File to colorize, by default stdin is used',
263 default='-',
264 nargs='?',
265 )
266 args = parser.parse_args()
267
268 if args.log == '-':
269 stream = sys.stdin
270 else:
271 stream = open(args.log, 'r')
272
273 if args.color == ColorArg.AUTO:
274 colors_enabled = sys.stdout.isatty()
275 elif args.color == ColorArg.ALWAYS:
276 colors_enabled = True
277 else:
278 colors_enabled = False
279
280 with contextlib.closing(stream):
281 colorize(stream, verbose=args.verbose, colors_enabled=colors_enabled)
282
283
284if __name__ == '__main__':
285 colorize_main()