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