1# Copyright 2015-2021 Espressif Systems (Shanghai) CO LTD
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import ctypes
16import os
17import re
18import sys
19from io import TextIOBase
20from typing import Any, Optional, TextIO, Union
21
22from .output_helpers import ANSI_NORMAL
23
24STD_OUTPUT_HANDLE = -11
25STD_ERROR_HANDLE = -12
26
27# wincon.h values
28FOREGROUND_INTENSITY = 8
29FOREGROUND_GREY = 7
30
31# matches the ANSI color change sequences that IDF sends
32RE_ANSI_COLOR = re.compile(b'\033\\[([01]);3([0-7])m')
33
34# list mapping the 8 ANSI colors (the indexes) to Windows Console colors
35ANSI_TO_WINDOWS_COLOR = [0, 4, 2, 6, 1, 5, 3, 7]
36
37if os.name == 'nt':
38    GetStdHandle = ctypes.windll.kernel32.GetStdHandle  # type: ignore
39    SetConsoleTextAttribute = ctypes.windll.kernel32.SetConsoleTextAttribute  # type: ignore
40
41
42def get_converter(orig_output_method=None, decode_output=False):
43    # type: (Any[TextIO, Optional[TextIOBase]], bool) -> Union[ANSIColorConverter, Optional[TextIOBase]]
44    """
45    Returns an ANSIColorConverter on Windows and the original output method (orig_output_method) on other platforms.
46    The ANSIColorConverter with decode_output=True will decode the bytes before passing them to the output.
47    """
48    if os.name == 'nt':
49        return ANSIColorConverter(orig_output_method, decode_output)
50    return orig_output_method
51
52
53class ANSIColorConverter(object):
54    """Class to wrap a file-like output stream, intercept ANSI color codes,
55    and convert them into calls to Windows SetConsoleTextAttribute.
56
57    Doesn't support all ANSI terminal code escape sequences, only the sequences IDF uses.
58
59    Ironically, in Windows this console output is normally wrapped by winpty which will then detect the console text
60    color changes and convert these back to ANSI color codes for MSYS' terminal to display. However this is the
61    least-bad working solution, as winpty doesn't support any "passthrough" mode for raw output.
62    """
63
64    def __init__(self, output=None, decode_output=False):
65        # type: (TextIOBase, bool) -> None
66        self.output = output
67        self.decode_output = decode_output
68        self.handle = GetStdHandle(STD_ERROR_HANDLE if self.output == sys.stderr else STD_OUTPUT_HANDLE)
69        self.matched = b''
70
71    def _output_write(self, data):  # type: (Union[str, bytes]) -> None
72        try:
73            if self.decode_output:
74                self.output.write(data.decode())  # type: ignore
75            else:
76                self.output.write(data)  # type: ignore
77        except (IOError, OSError):
78            # Windows 10 bug since the Fall Creators Update, sometimes writing to console randomly throws
79            # an exception (however, the character is still written to the screen)
80            # Ref https://github.com/espressif/esp-idf/issues/1163
81            #
82            # Also possible for Windows to throw an OSError error if the data is invalid for the console
83            # (garbage bytes, etc)
84            pass
85        except UnicodeDecodeError:
86            # In case of double byte Unicode characters display '?'
87            self.output.write('?')  # type: ignore
88
89    def write(self, data):  # type: ignore
90        if isinstance(data, bytes):
91            data = bytearray(data)
92        else:
93            data = bytearray(data, 'utf-8')
94        for b in data:
95            b = bytes([b])
96            length = len(self.matched)
97            if b == b'\033':  # ESC
98                self.matched = b
99            elif (length == 1 and b == b'[') or (1 < length < 7):
100                self.matched += b
101                if self.matched == ANSI_NORMAL.encode('latin-1'):  # reset console
102                    # Flush is required only with Python3 - switching color before it is printed would mess up the console
103                    self.flush()
104                    SetConsoleTextAttribute(self.handle, FOREGROUND_GREY)
105                    self.matched = b''
106                elif len(self.matched) == 7:  # could be an ANSI sequence
107                    m = re.match(RE_ANSI_COLOR, self.matched)
108                    if m is not None:
109                        color = ANSI_TO_WINDOWS_COLOR[int(m.group(2))]
110                        if m.group(1) == b'1':
111                            color |= FOREGROUND_INTENSITY
112                        # Flush is required only with Python3 - switching color before it is printed would mess up the console
113                        self.flush()
114                        SetConsoleTextAttribute(self.handle, color)
115                    else:
116                        self._output_write(self.matched)  # not an ANSI color code, display verbatim
117                    self.matched = b''
118            else:
119                self._output_write(b)
120                self.matched = b''
121
122    def flush(self):  # type: () -> None
123        try:
124            self.output.flush()  # type: ignore
125        except OSError:
126            # Account for Windows Console refusing to accept garbage bytes (serial noise, etc)
127            pass
128