userver: /data/code/service_template/third_party/userver/testsuite/pytest_plugins/pytest_userver/utils/colorize.py Source File
⚠️ This is the documentation for an old userver version. Click here to switch to the latest version.
All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Modules Pages Concepts
colorize.py
1import itertools
2import json
3import sys
4
5from . import tskv
6
7
8# Info about colors https://misc.flogisoft.com/bash/tip_colors_and_formatting
9class Colors:
10 BLACK = '\033[30m'
11 RED = '\033[31m'
12 GREEN = '\033[32m'
13 YELLOW = '\033[33m'
14 BLUE = '\033[34m'
15 GRAY = '\033[37m'
16 DARK_GRAY = '\033[37m'
17 BRIGHT_RED = '\033[91m'
18 BRIGHT_GREEN = '\033[92m'
19 BRIGHT_YELLOW = '\033[93m'
20 DEFAULT = '\033[0m'
21 DEFAULT_BG = '\033[49m'
22 BG_BLACK = '\033[40m'
23
24 # No bright red color, no close colors, no too dark colors
25 __NICE_COLORS = [
26 '\033[38;5;{}m'.format(x)
27 for x in itertools.chain(
28 range(2, 7),
29 range(10, 15),
30 range(38, 51, 2),
31 range(75, 87, 2),
32 range(112, 123, 3),
33 range(128, 159, 3),
34 range(164, 195, 3),
35 range(203, 231, 3),
36 )
37 ]
38
39 @staticmethod
40 def colorize(value):
41 ind = hash(value) % len(Colors.__NICE_COLORS)
42 return Colors.__NICE_COLORS[ind]
43
44
45LEVEL_COLORS = {
46 'TRACE': Colors.DARK_GRAY,
47 'DEBUG': Colors.GRAY,
48 'INFO': Colors.GREEN,
49 'WARNING': Colors.YELLOW,
50 'ERROR': Colors.RED,
51 'CRITICAL': Colors.BRIGHT_RED,
52 'none': Colors.DEFAULT,
53}
54HTTP_STATUS_COLORS = {
55 '2': Colors.GREEN,
56 '3': Colors.GREEN,
57 '4': Colors.YELLOW,
58 '5': Colors.RED,
59}
60
61
63 def __init__(self, *, verbose=False, colors_enabled=True):
64 self._requests = {}
65 self.verbose = verbose
66 self.colors_enabled = colors_enabled
67
68 def colorize_line(self, line):
69 if not line.startswith('tskv\t'):
70 return line
71 return self.colorize_tskv(line)
72
73 def colorize_tskv(self, line):
74 row = tskv.parse_line(line)
75 return self.colorize_row(row)
76
77 def colorize_row(self, row):
78 row = row.copy()
79 flowid = '-'.join([row.get(key, '') for key in ('link', 'trace_id')])
80
81 entry_type = row.pop('_type', None)
82 link = row.get('link', None)
83 level = row.pop('level', 'none')
84 text = row.pop('text', '')
85
86 extra_fields = []
87 if entry_type == 'request':
88 self._requests[link] = self._build_request_info(row)
89 if 'body' in row:
90 extra_fields.append(
91 'request_body='
92 + self.textcolor(
93 try_reformat_json(row.pop('body')), Colors.YELLOW,
94 ),
95 )
96 elif entry_type == 'response':
97 if 'meta_code' in row:
98 status_code = row.pop('meta_code')
99 extra_fields.append(
100 self._http_status('meta_code', status_code),
101 )
102 if not text:
103 text = 'Response finished'
104 extra_fields.append(
105 'response_body='
106 + self.textcolor(
107 try_reformat_json(row.pop('body')), Colors.YELLOW,
108 ),
109 )
110 elif entry_type == 'mockserver_request':
111 text = 'Mockserver request finished'
112 if 'meta_code' in row:
113 status_code = str(row.pop('meta_code'))
114 extra_fields.append(
115 self._http_status('meta_code', status_code),
116 )
117 for key in ('method', 'url', 'status', 'exc_info', 'delay'):
118 value = row.pop(key, None)
119 if value:
120 extra_fields.append(f'{key}={value}')
121
122 if link in self._requests:
123 logid = f'[{self._requests[link]}]'
124 elif link is not None:
125 logid = f'[{link}]'
126 else:
127 logid = '<userver>'
128
129 level_color = LEVEL_COLORS.get(level)
130 flow_color = Colors.colorize(flowid)
131
132 fields = [
133 self.textcolor(f'{level:<8}', level_color),
134 self.textcolor(logid, flow_color),
135 ]
136 if text:
137 fields.append(text)
138 elif self.verbose:
139 fields.append('<NO TEXT>')
140 else:
141 return None
142
143 fields.extend(extra_fields)
144 if self.verbose:
145 fields.extend([f'{k}={v}' for k, v in row.items()])
146 return ' '.join(fields)
147
148 def textcolor(self, text, color):
149 if not self.colors_enabled:
150 return str(text)
151 return f'{color}{text}{Colors.DEFAULT}'
152
153 def _http_status(self, key, status):
154 color = HTTP_STATUS_COLORS.get(status[:1], Colors.DEFAULT)
155 return self.textcolor(f'{key}={status}', color)
156
157 def _build_request_info(self, row):
158 if 'uri' not in row:
159 return None
160 uri = row['uri']
161 method = row.get('method', 'UNKNOWN')
162 return f'{method} {uri}'
163
164
165def format_json(obj):
166 encoded = json.dumps(
167 obj,
168 indent=2,
169 separators=(',', ': '),
170 sort_keys=True,
171 ensure_ascii=False,
172 )
173 return encoded
174
175
176def try_reformat_json(body):
177 try:
178 # TODO: unescape string
179 data = json.loads(body)
180 return format_json(data)
181 except ValueError:
182 return body
183
184
185def colorize(stream):
186 colorizer = Colorizer()
187 for line in stream:
188 line = line.rstrip('\r\n')
189 color_line = colorizer.colorize_line(line)
190 if color_line is not None:
191 print(color_line)
192
193
194if __name__ == '__main__':
195 colorize(sys.stdin)