userver: /data/code/userver/testsuite/pytest_plugins/pytest_userver/utils/colorize.py Source File
Loading...
Searching...
No Matches
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 'meta_code' in row:
110 status_code = row.pop('meta_code')
111 extra_fields.append(
112 self._http_status('meta_code', status_code),
113 )
114 if not text:
115 text = 'Response finished'
116 extra_fields.append(
117 'response_body='
118 + self.textcolor(
119 try_reformat_json(row.pop('body')),
120 Colors.YELLOW,
121 ),
122 )
123 elif entry_type == 'mockserver_request':
124 text = 'Mockserver request finished'
125 if 'meta_code' in row:
126 status_code = str(row.pop('meta_code'))
127 extra_fields.append(
128 self._http_status('meta_code', status_code),
129 )
130 for key in ('method', 'url', 'status', 'exc_info', 'delay'):
131 value = row.pop(key, None)
132 if value:
133 extra_fields.append(f'{key}={value}')
134
135 if link in self._requests:
136 logid = f'[{self._requests[link]}]'
137 elif link is not None:
138 logid = f'[{link}]'
139 else:
140 logid = '<userver>'
141
142 level_color = LEVEL_COLORS.get(level)
143 flow_color = Colors.colorize(flowid)
144
145 fields = [
146 self.textcolor(f'{level:<8}', level_color),
147 self.textcolor(logid, flow_color),
148 ]
149 if text:
150 http_url_info = row.pop('http_url', None)
151 if http_url_info:
152 http_url_info = http_url_info.removeprefix(
153 HTTP_LOCALHOST_PREFIX,
154 )
155 http_url_info = http_url_info.removesuffix('?')
156 http_url_info = http_url_info[http_url_info.find('/') + 1 :]
157
158 meta_code = row.pop('meta_code', None)
159 if meta_code:
160 http_url_info += f' meta_code={meta_code}'
161
162 text = f'{self.textcolor(http_url_info, Colors.GREEN)} {text}'
163
164 fields.append(text)
165 elif self.verbose:
166 fields.append('<NO TEXT>')
167 else:
168 return None
169
170 fields.extend(extra_fields)
171 if self.verbose:
172 fields.extend([f'{k}={v}' for k, v in row.items()])
173 return ' '.join(fields)
174
175 def textcolor(self, text, color):
176 if not self.colors_enabled:
177 return str(text)
178 return f'{color}{text}{Colors.DEFAULT}'
179
180 def _http_status(self, key, status):
181 color = HTTP_STATUS_COLORS.get(status[:1], Colors.DEFAULT)
182 return self.textcolor(f'{key}={status}', color)
183
184 def _build_request_info(self, row):
185 if 'uri' not in row:
186 return None
187 uri = row['uri']
188 method = row.get('method', 'UNKNOWN')
189 return f'{method} {uri}'
190
191
192def format_json(obj):
193 encoded = json.dumps(
194 obj,
195 indent=2,
196 separators=(',', ': '),
197 sort_keys=True,
198 ensure_ascii=False,
199 )
200 return encoded
201
202
203def try_reformat_json(body):
204 try:
205 # TODO: unescape string
206 data = json.loads(body)
207 return format_json(data)
208 except ValueError:
209 return body
210
211
212def colorize(stream, verbose=False, colors_enabled=True):
213 colorizer = Colorizer(verbose=verbose, colors_enabled=colors_enabled)
214 for line in stream:
215 line = line.rstrip('\r\n')
216 color_line = colorizer.colorize_line(line)
217 if color_line is not None:
218 print(color_line)
219
220
221def parse_color(value):
222 if value in ('always', 'force', 'yes', 'enable'):
223 return ColorArg.ALWAYS
224 if value in ('never', 'no', 'disable'):
225 return ColorArg.NEVER
226 if value == 'auto':
227 return ColorArg.AUTO
228 raise ValueError(f'Unknown color option {value!r}')
229
230
231def colorize_main():
232 parser = argparse.ArgumentParser(description='Colorize userver log file.')
233 parser.add_argument(
234 '--verbose',
235 '-v',
236 action='store_true',
237 help='Be verbose',
238 )
239 parser.add_argument(
240 '--color',
241 metavar='WHEN',
242 help=('Control color highlighting, WHEN is always, never or ' 'auto (default)'),
243 nargs='?',
244 type=parse_color,
245 default=ColorArg.AUTO,
246 const=ColorArg.ALWAYS,
247 )
248 parser.add_argument(
249 '--no-color',
250 help='Turn off color highlighting',
251 dest='color',
252 action='store_const',
253 const=ColorArg.NEVER,
254 )
255 parser.add_argument(
256 'log',
257 help='File to colorize, by default stdin is used',
258 default='-',
259 nargs='?',
260 )
261 args = parser.parse_args()
262
263 if args.log == '-':
264 stream = sys.stdin
265 else:
266 stream = open(args.log, 'r')
267
268 if args.color == ColorArg.AUTO:
269 colors_enabled = sys.stdout.isatty()
270 elif args.color == ColorArg.ALWAYS:
271 colors_enabled = True
272 else:
273 colors_enabled = False
274
275 with contextlib.closing(stream):
276 colorize(stream, verbose=args.verbose, colors_enabled=colors_enabled)
277
278
279if __name__ == '__main__':
280 colorize_main()