1# DExTer : Debugging Experience Tester 2# ~~~~~~ ~ ~~ ~ ~~ 3# 4# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 5# See https://llvm.org/LICENSE.txt for license information. 6# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 7"""Provides formatted/colored console output on both Windows and Linux. 8 9Do not use this module directly, but instead use via the appropriate platform- 10specific module. 11""" 12 13import abc 14import re 15import sys 16import threading 17import unittest 18 19from io import StringIO 20 21from dex.utils.Exceptions import Error 22 23 24class _NullLock(object): 25 def __enter__(self): 26 return None 27 28 def __exit__(self, *params): 29 pass 30 31 32_lock = threading.Lock() 33_null_lock = _NullLock() 34 35 36class PreserveAutoColors(object): 37 def __init__(self, pretty_output): 38 self.pretty_output = pretty_output 39 self.orig_values = {} 40 self.properties = [ 41 'auto_reds', 'auto_yellows', 'auto_greens', 'auto_blues' 42 ] 43 44 def __enter__(self): 45 for p in self.properties: 46 self.orig_values[p] = getattr(self.pretty_output, p)[:] 47 return self 48 49 def __exit__(self, *args): 50 for p in self.properties: 51 setattr(self.pretty_output, p, self.orig_values[p]) 52 53 54class Stream(object): 55 def __init__(self, py_, os_=None): 56 self.py = py_ 57 self.os = os_ 58 self.orig_color = None 59 self.color_enabled = self.py.isatty() 60 61 62class PrettyOutputBase(object, metaclass=abc.ABCMeta): 63 stdout = Stream(sys.stdout) 64 stderr = Stream(sys.stderr) 65 66 def __init__(self): 67 self.auto_reds = [] 68 self.auto_yellows = [] 69 self.auto_greens = [] 70 self.auto_blues = [] 71 self._stack = [] 72 73 def __enter__(self): 74 return self 75 76 def __exit__(self, *args): 77 pass 78 79 def _set_valid_stream(self, stream): 80 if stream is None: 81 return self.__class__.stdout 82 return stream 83 84 def _write(self, text, stream): 85 text = str(text) 86 87 # Users can embed color control tags in their output 88 # (e.g. <r>hello</> <y>world</> would write the word 'hello' in red and 89 # 'world' in yellow). 90 # This function parses these tags using a very simple recursive 91 # descent. 92 colors = { 93 'r': self.red, 94 'y': self.yellow, 95 'g': self.green, 96 'b': self.blue, 97 'd': self.default, 98 'a': self.auto, 99 } 100 101 # Find all tags (whether open or close) 102 tags = [ 103 t for t in re.finditer('<([{}/])>'.format(''.join(colors)), text) 104 ] 105 106 if not tags: 107 # No tags. Just write the text to the current stream and return. 108 # 'unmangling' any tags that have been mangled so that they won't 109 # render as colors (for example in error output from this 110 # function). 111 stream = self._set_valid_stream(stream) 112 stream.py.write(text.replace(r'\>', '>')) 113 return 114 115 open_tags = [i for i in tags if i.group(1) != '/'] 116 close_tags = [i for i in tags if i.group(1) == '/'] 117 118 if (len(open_tags) != len(close_tags) 119 or any(o.start() >= c.start() 120 for (o, c) in zip(open_tags, close_tags))): 121 raise Error('open/close tag mismatch in "{}"'.format( 122 text.rstrip()).replace('>', r'\>')) 123 124 open_tag = open_tags.pop(0) 125 126 # We know that the tags balance correctly, so figure out where the 127 # corresponding close tag is to the current open tag. 128 tag_nesting = 1 129 close_tag = None 130 for tag in tags[1:]: 131 if tag.group(1) == '/': 132 tag_nesting -= 1 133 else: 134 tag_nesting += 1 135 if tag_nesting == 0: 136 close_tag = tag 137 break 138 else: 139 assert False, text 140 141 # Use the method on the top of the stack for text prior to the open 142 # tag. 143 before = text[:open_tag.start()] 144 if before: 145 self._stack[-1](before, lock=_null_lock, stream=stream) 146 147 # Use the specified color for the tag itself. 148 color = open_tag.group(1) 149 within = text[open_tag.end():close_tag.start()] 150 if within: 151 colors[color](within, lock=_null_lock, stream=stream) 152 153 # Use the method on the top of the stack for text after the close tag. 154 after = text[close_tag.end():] 155 if after: 156 self._stack[-1](after, lock=_null_lock, stream=stream) 157 158 def flush(self, stream): 159 stream = self._set_valid_stream(stream) 160 stream.py.flush() 161 162 def auto(self, text, stream=None, lock=_lock): 163 text = str(text) 164 stream = self._set_valid_stream(stream) 165 lines = text.splitlines(True) 166 167 with lock: 168 for line in lines: 169 # This is just being cute for the sake of cuteness, but why 170 # not? 171 line = line.replace('DExTer', '<r>D<y>E<g>x<b>T</></>e</>r</>') 172 173 # Apply the appropriate color method if the expression matches 174 # any of 175 # the patterns we have set up. 176 for fn, regexs in ((self.red, self.auto_reds), 177 (self.yellow, self.auto_yellows), 178 (self.green, 179 self.auto_greens), (self.blue, 180 self.auto_blues)): 181 if any(re.search(regex, line) for regex in regexs): 182 fn(line, stream=stream, lock=_null_lock) 183 break 184 else: 185 self.default(line, stream=stream, lock=_null_lock) 186 187 def _call_color_impl(self, fn, impl, text, *args, **kwargs): 188 try: 189 self._stack.append(fn) 190 return impl(text, *args, **kwargs) 191 finally: 192 fn = self._stack.pop() 193 194 @abc.abstractmethod 195 def red_impl(self, text, stream=None, **kwargs): 196 pass 197 198 def red(self, *args, **kwargs): 199 return self._call_color_impl(self.red, self.red_impl, *args, **kwargs) 200 201 @abc.abstractmethod 202 def yellow_impl(self, text, stream=None, **kwargs): 203 pass 204 205 def yellow(self, *args, **kwargs): 206 return self._call_color_impl(self.yellow, self.yellow_impl, *args, 207 **kwargs) 208 209 @abc.abstractmethod 210 def green_impl(self, text, stream=None, **kwargs): 211 pass 212 213 def green(self, *args, **kwargs): 214 return self._call_color_impl(self.green, self.green_impl, *args, 215 **kwargs) 216 217 @abc.abstractmethod 218 def blue_impl(self, text, stream=None, **kwargs): 219 pass 220 221 def blue(self, *args, **kwargs): 222 return self._call_color_impl(self.blue, self.blue_impl, *args, 223 **kwargs) 224 225 @abc.abstractmethod 226 def default_impl(self, text, stream=None, **kwargs): 227 pass 228 229 def default(self, *args, **kwargs): 230 return self._call_color_impl(self.default, self.default_impl, *args, 231 **kwargs) 232 233 def colortest(self): 234 from itertools import combinations, permutations 235 236 fns = ((self.red, 'rrr'), (self.yellow, 'yyy'), (self.green, 'ggg'), 237 (self.blue, 'bbb'), (self.default, 'ddd')) 238 239 for l in range(1, len(fns) + 1): 240 for comb in combinations(fns, l): 241 for perm in permutations(comb): 242 for stream in (None, self.__class__.stderr): 243 perm[0][0]('stdout ' 244 if stream is None else 'stderr ', stream) 245 for fn, string in perm: 246 fn(string, stream) 247 self.default('\n', stream) 248 249 tests = [ 250 (self.auto, 'default1<r>red2</>default3'), 251 (self.red, 'red1<r>red2</>red3'), 252 (self.blue, 'blue1<r>red2</>blue3'), 253 (self.red, 'red1<y>yellow2</>red3'), 254 (self.auto, 'default1<y>yellow2<r>red3</></>'), 255 (self.auto, 'default1<g>green2<r>red3</></>'), 256 (self.auto, 'default1<g>green2<r>red3</>green4</>default5'), 257 (self.auto, 'default1<g>green2</>default3<g>green4</>default5'), 258 (self.auto, '<r>red1<g>green2</>red3<g>green4</>red5</>'), 259 (self.auto, '<r>red1<y><g>green2</>yellow3</>green4</>default5'), 260 (self.auto, '<r><y><g><b><d>default1</></><r></></></>red2</>'), 261 (self.auto, '<r>red1</>default2<r>red3</><g>green4</>default5'), 262 (self.blue, '<r>red1</>blue2<r><r>red3</><g><g>green</></></>'), 263 (self.blue, '<r>r<r>r<y>y<r><r><r><r>r</></></></></></></>b'), 264 ] 265 266 for fn, text in tests: 267 for stream in (None, self.__class__.stderr): 268 stream_name = 'stdout' if stream is None else 'stderr' 269 fn('{} {}\n'.format(stream_name, text), stream) 270 271 272class TestPrettyOutput(unittest.TestCase): 273 class MockPrettyOutput(PrettyOutputBase): 274 def red_impl(self, text, stream=None, **kwargs): 275 self._write('[R]{}[/R]'.format(text), stream) 276 277 def yellow_impl(self, text, stream=None, **kwargs): 278 self._write('[Y]{}[/Y]'.format(text), stream) 279 280 def green_impl(self, text, stream=None, **kwargs): 281 self._write('[G]{}[/G]'.format(text), stream) 282 283 def blue_impl(self, text, stream=None, **kwargs): 284 self._write('[B]{}[/B]'.format(text), stream) 285 286 def default_impl(self, text, stream=None, **kwargs): 287 self._write('[D]{}[/D]'.format(text), stream) 288 289 def test_red(self): 290 with TestPrettyOutput.MockPrettyOutput() as o: 291 stream = Stream(StringIO()) 292 o.red('hello', stream) 293 self.assertEqual(stream.py.getvalue(), '[R]hello[/R]') 294 295 def test_yellow(self): 296 with TestPrettyOutput.MockPrettyOutput() as o: 297 stream = Stream(StringIO()) 298 o.yellow('hello', stream) 299 self.assertEqual(stream.py.getvalue(), '[Y]hello[/Y]') 300 301 def test_green(self): 302 with TestPrettyOutput.MockPrettyOutput() as o: 303 stream = Stream(StringIO()) 304 o.green('hello', stream) 305 self.assertEqual(stream.py.getvalue(), '[G]hello[/G]') 306 307 def test_blue(self): 308 with TestPrettyOutput.MockPrettyOutput() as o: 309 stream = Stream(StringIO()) 310 o.blue('hello', stream) 311 self.assertEqual(stream.py.getvalue(), '[B]hello[/B]') 312 313 def test_default(self): 314 with TestPrettyOutput.MockPrettyOutput() as o: 315 stream = Stream(StringIO()) 316 o.default('hello', stream) 317 self.assertEqual(stream.py.getvalue(), '[D]hello[/D]') 318 319 def test_auto(self): 320 with TestPrettyOutput.MockPrettyOutput() as o: 321 stream = Stream(StringIO()) 322 o.auto_reds.append('foo') 323 o.auto('bar\n', stream) 324 o.auto('foo\n', stream) 325 o.auto('baz\n', stream) 326 self.assertEqual(stream.py.getvalue(), 327 '[D]bar\n[/D][R]foo\n[/R][D]baz\n[/D]') 328 329 stream = Stream(StringIO()) 330 o.auto('bar\nfoo\nbaz\n', stream) 331 self.assertEqual(stream.py.getvalue(), 332 '[D]bar\n[/D][R]foo\n[/R][D]baz\n[/D]') 333 334 stream = Stream(StringIO()) 335 o.auto('barfoobaz\nbardoobaz\n', stream) 336 self.assertEqual(stream.py.getvalue(), 337 '[R]barfoobaz\n[/R][D]bardoobaz\n[/D]') 338 339 o.auto_greens.append('doo') 340 stream = Stream(StringIO()) 341 o.auto('barfoobaz\nbardoobaz\n', stream) 342 self.assertEqual(stream.py.getvalue(), 343 '[R]barfoobaz\n[/R][G]bardoobaz\n[/G]') 344 345 def test_PreserveAutoColors(self): 346 with TestPrettyOutput.MockPrettyOutput() as o: 347 o.auto_reds.append('foo') 348 with PreserveAutoColors(o): 349 o.auto_greens.append('bar') 350 stream = Stream(StringIO()) 351 o.auto('foo\nbar\nbaz\n', stream) 352 self.assertEqual(stream.py.getvalue(), 353 '[R]foo\n[/R][G]bar\n[/G][D]baz\n[/D]') 354 355 stream = Stream(StringIO()) 356 o.auto('foo\nbar\nbaz\n', stream) 357 self.assertEqual(stream.py.getvalue(), 358 '[R]foo\n[/R][D]bar\n[/D][D]baz\n[/D]') 359 360 stream = Stream(StringIO()) 361 o.yellow('<a>foo</>bar<a>baz</>', stream) 362 self.assertEqual( 363 stream.py.getvalue(), 364 '[Y][Y][/Y][R]foo[/R][Y][Y]bar[/Y][D]baz[/D][Y][/Y][/Y][/Y]') 365 366 def test_tags(self): 367 with TestPrettyOutput.MockPrettyOutput() as o: 368 stream = Stream(StringIO()) 369 o.auto('<r>hi</>', stream) 370 self.assertEqual(stream.py.getvalue(), 371 '[D][D][/D][R]hi[/R][D][/D][/D]') 372 373 stream = Stream(StringIO()) 374 o.auto('<r><y>a</>b</>c', stream) 375 self.assertEqual( 376 stream.py.getvalue(), 377 '[D][D][/D][R][R][/R][Y]a[/Y][R]b[/R][/R][D]c[/D][/D]') 378 379 with self.assertRaisesRegex(Error, 'tag mismatch'): 380 o.auto('<r>hi', stream) 381 382 with self.assertRaisesRegex(Error, 'tag mismatch'): 383 o.auto('hi</>', stream) 384 385 with self.assertRaisesRegex(Error, 'tag mismatch'): 386 o.auto('<r><y>hi</>', stream) 387 388 with self.assertRaisesRegex(Error, 'tag mismatch'): 389 o.auto('<r><y>hi</><r></>', stream) 390 391 with self.assertRaisesRegex(Error, 'tag mismatch'): 392 o.auto('</>hi<r>', stream) 393