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 = ["auto_reds", "auto_yellows", "auto_greens", "auto_blues"] 41 42 def __enter__(self): 43 for p in self.properties: 44 self.orig_values[p] = getattr(self.pretty_output, p)[:] 45 return self 46 47 def __exit__(self, *args): 48 for p in self.properties: 49 setattr(self.pretty_output, p, self.orig_values[p]) 50 51 52class Stream(object): 53 def __init__(self, py_, os_=None): 54 self.py = py_ 55 self.os = os_ 56 self.orig_color = None 57 self.color_enabled = self.py.isatty() 58 59 60class PrettyOutputBase(object, metaclass=abc.ABCMeta): 61 stdout = Stream(sys.stdout) 62 stderr = Stream(sys.stderr) 63 64 def __init__(self): 65 self.auto_reds = [] 66 self.auto_yellows = [] 67 self.auto_greens = [] 68 self.auto_blues = [] 69 self._stack = [] 70 71 def __enter__(self): 72 return self 73 74 def __exit__(self, *args): 75 pass 76 77 def _set_valid_stream(self, stream): 78 if stream is None: 79 return self.__class__.stdout 80 return stream 81 82 def _write(self, text, stream): 83 text = str(text) 84 85 # Users can embed color control tags in their output 86 # (e.g. <r>hello</> <y>world</> would write the word 'hello' in red and 87 # 'world' in yellow). 88 # This function parses these tags using a very simple recursive 89 # descent. 90 colors = { 91 "r": self.red, 92 "y": self.yellow, 93 "g": self.green, 94 "b": self.blue, 95 "d": self.default, 96 "a": self.auto, 97 } 98 99 # Find all tags (whether open or close) 100 tags = [t for t in re.finditer("<([{}/])>".format("".join(colors)), text)] 101 102 if not tags: 103 # No tags. Just write the text to the current stream and return. 104 # 'unmangling' any tags that have been mangled so that they won't 105 # render as colors (for example in error output from this 106 # function). 107 stream = self._set_valid_stream(stream) 108 stream.py.write(text.replace(r"\>", ">")) 109 return 110 111 open_tags = [i for i in tags if i.group(1) != "/"] 112 close_tags = [i for i in tags if i.group(1) == "/"] 113 114 if len(open_tags) != len(close_tags) or any( 115 o.start() >= c.start() for (o, c) in zip(open_tags, close_tags) 116 ): 117 raise Error( 118 'open/close tag mismatch in "{}"'.format(text.rstrip()).replace( 119 ">", r"\>" 120 ) 121 ) 122 123 open_tag = open_tags.pop(0) 124 125 # We know that the tags balance correctly, so figure out where the 126 # corresponding close tag is to the current open tag. 127 tag_nesting = 1 128 close_tag = None 129 for tag in tags[1:]: 130 if tag.group(1) == "/": 131 tag_nesting -= 1 132 else: 133 tag_nesting += 1 134 if tag_nesting == 0: 135 close_tag = tag 136 break 137 else: 138 assert False, text 139 140 # Use the method on the top of the stack for text prior to the open 141 # tag. 142 before = text[: open_tag.start()] 143 if before: 144 self._stack[-1](before, lock=_null_lock, stream=stream) 145 146 # Use the specified color for the tag itself. 147 color = open_tag.group(1) 148 within = text[open_tag.end() : close_tag.start()] 149 if within: 150 colors[color](within, lock=_null_lock, stream=stream) 151 152 # Use the method on the top of the stack for text after the close tag. 153 after = text[close_tag.end() :] 154 if after: 155 self._stack[-1](after, lock=_null_lock, stream=stream) 156 157 def flush(self, stream): 158 stream = self._set_valid_stream(stream) 159 stream.py.flush() 160 161 def auto(self, text, stream=None, lock=_lock): 162 text = str(text) 163 stream = self._set_valid_stream(stream) 164 lines = text.splitlines(True) 165 166 with lock: 167 for line in lines: 168 # This is just being cute for the sake of cuteness, but why 169 # not? 170 line = line.replace("DExTer", "<r>D<y>E<g>x<b>T</></>e</>r</>") 171 172 # Apply the appropriate color method if the expression matches 173 # any of 174 # the patterns we have set up. 175 for fn, regexs in ( 176 (self.red, self.auto_reds), 177 (self.yellow, self.auto_yellows), 178 (self.green, self.auto_greens), 179 (self.blue, self.auto_blues), 180 ): 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, **kwargs) 207 208 @abc.abstractmethod 209 def green_impl(self, text, stream=None, **kwargs): 210 pass 211 212 def green(self, *args, **kwargs): 213 return self._call_color_impl(self.green, self.green_impl, *args, **kwargs) 214 215 @abc.abstractmethod 216 def blue_impl(self, text, stream=None, **kwargs): 217 pass 218 219 def blue(self, *args, **kwargs): 220 return self._call_color_impl(self.blue, self.blue_impl, *args, **kwargs) 221 222 @abc.abstractmethod 223 def default_impl(self, text, stream=None, **kwargs): 224 pass 225 226 def default(self, *args, **kwargs): 227 return self._call_color_impl(self.default, self.default_impl, *args, **kwargs) 228 229 def colortest(self): 230 from itertools import combinations, permutations 231 232 fns = ( 233 (self.red, "rrr"), 234 (self.yellow, "yyy"), 235 (self.green, "ggg"), 236 (self.blue, "bbb"), 237 (self.default, "ddd"), 238 ) 239 240 for l in range(1, len(fns) + 1): 241 for comb in combinations(fns, l): 242 for perm in permutations(comb): 243 for stream in (None, self.__class__.stderr): 244 perm[0][0]("stdout " 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( 327 stream.py.getvalue(), "[D]bar\n[/D][R]foo\n[/R][D]baz\n[/D]" 328 ) 329 330 stream = Stream(StringIO()) 331 o.auto("bar\nfoo\nbaz\n", stream) 332 self.assertEqual( 333 stream.py.getvalue(), "[D]bar\n[/D][R]foo\n[/R][D]baz\n[/D]" 334 ) 335 336 stream = Stream(StringIO()) 337 o.auto("barfoobaz\nbardoobaz\n", stream) 338 self.assertEqual( 339 stream.py.getvalue(), "[R]barfoobaz\n[/R][D]bardoobaz\n[/D]" 340 ) 341 342 o.auto_greens.append("doo") 343 stream = Stream(StringIO()) 344 o.auto("barfoobaz\nbardoobaz\n", stream) 345 self.assertEqual( 346 stream.py.getvalue(), "[R]barfoobaz\n[/R][G]bardoobaz\n[/G]" 347 ) 348 349 def test_PreserveAutoColors(self): 350 with TestPrettyOutput.MockPrettyOutput() as o: 351 o.auto_reds.append("foo") 352 with PreserveAutoColors(o): 353 o.auto_greens.append("bar") 354 stream = Stream(StringIO()) 355 o.auto("foo\nbar\nbaz\n", stream) 356 self.assertEqual( 357 stream.py.getvalue(), "[R]foo\n[/R][G]bar\n[/G][D]baz\n[/D]" 358 ) 359 360 stream = Stream(StringIO()) 361 o.auto("foo\nbar\nbaz\n", stream) 362 self.assertEqual( 363 stream.py.getvalue(), "[R]foo\n[/R][D]bar\n[/D][D]baz\n[/D]" 364 ) 365 366 stream = Stream(StringIO()) 367 o.yellow("<a>foo</>bar<a>baz</>", stream) 368 self.assertEqual( 369 stream.py.getvalue(), 370 "[Y][Y][/Y][R]foo[/R][Y][Y]bar[/Y][D]baz[/D][Y][/Y][/Y][/Y]", 371 ) 372 373 def test_tags(self): 374 with TestPrettyOutput.MockPrettyOutput() as o: 375 stream = Stream(StringIO()) 376 o.auto("<r>hi</>", stream) 377 self.assertEqual(stream.py.getvalue(), "[D][D][/D][R]hi[/R][D][/D][/D]") 378 379 stream = Stream(StringIO()) 380 o.auto("<r><y>a</>b</>c", stream) 381 self.assertEqual( 382 stream.py.getvalue(), 383 "[D][D][/D][R][R][/R][Y]a[/Y][R]b[/R][/R][D]c[/D][/D]", 384 ) 385 386 with self.assertRaisesRegex(Error, "tag mismatch"): 387 o.auto("<r>hi", stream) 388 389 with self.assertRaisesRegex(Error, "tag mismatch"): 390 o.auto("hi</>", stream) 391 392 with self.assertRaisesRegex(Error, "tag mismatch"): 393 o.auto("<r><y>hi</>", stream) 394 395 with self.assertRaisesRegex(Error, "tag mismatch"): 396 o.auto("<r><y>hi</><r></>", stream) 397 398 with self.assertRaisesRegex(Error, "tag mismatch"): 399 o.auto("</>hi<r>", stream) 400