400 lines · python
1# DExTer : Debugging Experience Tester2# ~~~~~~ ~ ~~ ~ ~~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-exception7"""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 abc14import re15import sys16import threading17import unittest18 19from io import StringIO20 21from dex.utils.Exceptions import Error22 23 24class _NullLock(object):25 def __enter__(self):26 return None27 28 def __exit__(self, *params):29 pass30 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_output39 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 self46 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 = None57 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 self73 74 def __exit__(self, *args):75 pass76 77 def _set_valid_stream(self, stream):78 if stream is None:79 return self.__class__.stdout80 return stream81 82 def _write(self, text, stream):83 text = str(text)84 85 # Users can embed color control tags in their output86 # (e.g. <r>hello</> <y>world</> would write the word 'hello' in red and87 # 'world' in yellow).88 # This function parses these tags using a very simple recursive89 # 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't105 # render as colors (for example in error output from this106 # function).107 stream = self._set_valid_stream(stream)108 stream.py.write(text.replace(r"\>", ">"))109 return110 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 the126 # corresponding close tag is to the current open tag.127 tag_nesting = 1128 close_tag = None129 for tag in tags[1:]:130 if tag.group(1) == "/":131 tag_nesting -= 1132 else:133 tag_nesting += 1134 if tag_nesting == 0:135 close_tag = tag136 break137 else:138 assert False, text139 140 # Use the method on the top of the stack for text prior to the open141 # 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 why169 # 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 matches173 # any of174 # 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 break184 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.abstractmethod195 def red_impl(self, text, stream=None, **kwargs):196 pass197 198 def red(self, *args, **kwargs):199 return self._call_color_impl(self.red, self.red_impl, *args, **kwargs)200 201 @abc.abstractmethod202 def yellow_impl(self, text, stream=None, **kwargs):203 pass204 205 def yellow(self, *args, **kwargs):206 return self._call_color_impl(self.yellow, self.yellow_impl, *args, **kwargs)207 208 @abc.abstractmethod209 def green_impl(self, text, stream=None, **kwargs):210 pass211 212 def green(self, *args, **kwargs):213 return self._call_color_impl(self.green, self.green_impl, *args, **kwargs)214 215 @abc.abstractmethod216 def blue_impl(self, text, stream=None, **kwargs):217 pass218 219 def blue(self, *args, **kwargs):220 return self._call_color_impl(self.blue, self.blue_impl, *args, **kwargs)221 222 @abc.abstractmethod223 def default_impl(self, text, stream=None, **kwargs):224 pass225 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, permutations231 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