efro.terminal

  1# Released under the MIT License. See LICENSE for details.
  2#
  3"""Functionality related to terminal IO."""
  4from __future__ import annotations
  5
  6import sys
  7import os
  8from enum import Enum, unique
  9from typing import TYPE_CHECKING
 10
 11if TYPE_CHECKING:
 12    from typing import Any, ClassVar
 13
 14
 15@unique
 16class TerminalColor(Enum):
 17    """Color codes for printing to terminals.
 18
 19    Generally the Clr class should be used when incorporating color into
 20    terminal output, as it handles non-color-supporting terminals/etc.
 21    """
 22
 23    # Styles
 24    RESET = '\033[0m'
 25    BOLD = '\033[1m'
 26    UNDERLINE = '\033[4m'
 27    INVERSE = '\033[7m'
 28
 29    # Normal foreground colors
 30    BLACK = '\033[30m'
 31    RED = '\033[31m'
 32    GREEN = '\033[32m'
 33    YELLOW = '\033[33m'
 34    BLUE = '\033[34m'
 35    MAGENTA = '\033[35m'
 36    CYAN = '\033[36m'
 37    WHITE = '\033[37m'
 38
 39    # Normal background colors.
 40    BG_BLACK = '\033[40m'
 41    BG_RED = '\033[41m'
 42    BG_GREEN = '\033[42m'
 43    BG_YELLOW = '\033[43m'
 44    BG_BLUE = '\033[44m'
 45    BG_MAGENTA = '\033[45m'
 46    BG_CYAN = '\033[46m'
 47    BG_WHITE = '\033[47m'
 48
 49    # Strong foreground colors
 50    STRONG_BLACK = '\033[90m'
 51    STRONG_RED = '\033[91m'
 52    STRONG_GREEN = '\033[92m'
 53    STRONG_YELLOW = '\033[93m'
 54    STRONG_BLUE = '\033[94m'
 55    STRONG_MAGENTA = '\033[95m'
 56    STRONG_CYAN = '\033[96m'
 57    STRONG_WHITE = '\033[97m'
 58
 59    # Strong background colors.
 60    STRONG_BG_BLACK = '\033[100m'
 61    STRONG_BG_RED = '\033[101m'
 62    STRONG_BG_GREEN = '\033[102m'
 63    STRONG_BG_YELLOW = '\033[103m'
 64    STRONG_BG_BLUE = '\033[104m'
 65    STRONG_BG_MAGENTA = '\033[105m'
 66    STRONG_BG_CYAN = '\033[106m'
 67    STRONG_BG_WHITE = '\033[107m'
 68
 69
 70def _default_color_enabled() -> bool:
 71    """Return whether we enable ANSI color codes by default."""
 72    import platform
 73
 74    # If our stdout is not attached to a terminal, go with no-color.
 75    assert sys.__stdout__ is not None
 76    if not sys.__stdout__.isatty():
 77        return False
 78
 79    termenv = os.environ.get('TERM')
 80
 81    # If TERM is unset, don't attempt color (this is currently the case
 82    # in xcode).
 83    if termenv is None:
 84        return False
 85
 86    # A common way to say the terminal can't do fancy stuff like color.
 87    if termenv == 'dumb':
 88        return False
 89
 90    # On windows, try to enable ANSI color mode.
 91    if platform.system() == 'Windows':
 92        return _windows_enable_color()
 93
 94    # We seem to be a terminal with color support; let's do it!
 95    return True
 96
 97
 98# noinspection PyPep8Naming
 99def _windows_enable_color() -> bool:
100    """Attempt to enable ANSI color on windows terminal; return success."""
101    # pylint: disable=invalid-name, import-error, undefined-variable
102    # Pulled from: https://bugs.python.org/issue30075
103    import msvcrt
104    import ctypes
105    from ctypes import wintypes
106
107    kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)  # type: ignore
108
109    ERROR_INVALID_PARAMETER = 0x0057
110    ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004
111
112    def _check_bool(result: Any, _func: Any, args: Any) -> Any:
113        if not result:
114            raise ctypes.WinError(ctypes.get_last_error())  # type: ignore
115        return args
116
117    LPDWORD = ctypes.POINTER(wintypes.DWORD)
118    kernel32.GetConsoleMode.errcheck = _check_bool
119    kernel32.GetConsoleMode.argtypes = (wintypes.HANDLE, LPDWORD)
120    kernel32.SetConsoleMode.errcheck = _check_bool
121    kernel32.SetConsoleMode.argtypes = (wintypes.HANDLE, wintypes.DWORD)
122
123    def set_conout_mode(new_mode: int, mask: int = 0xFFFFFFFF) -> int:
124        # don't assume StandardOutput is a console.
125        # open CONOUT$ instead
126        fdout = os.open('CONOUT$', os.O_RDWR)
127        try:
128            hout = msvcrt.get_osfhandle(fdout)  # type: ignore
129            # pylint: disable=useless-suppression
130            # pylint: disable=no-value-for-parameter
131            old_mode = wintypes.DWORD()
132            # pylint: enable=useless-suppression
133            kernel32.GetConsoleMode(hout, ctypes.byref(old_mode))
134            mode = (new_mode & mask) | (old_mode.value & ~mask)
135            kernel32.SetConsoleMode(hout, mode)
136            return old_mode.value
137        finally:
138            os.close(fdout)
139
140    def enable_vt_mode() -> int:
141        mode = mask = ENABLE_VIRTUAL_TERMINAL_PROCESSING
142        try:
143            return set_conout_mode(mode, mask)
144        except WindowsError as exc:  # type: ignore
145            if exc.winerror == ERROR_INVALID_PARAMETER:
146                raise NotImplementedError from exc
147            raise
148
149    try:
150        enable_vt_mode()
151        return True
152    except NotImplementedError:
153        return False
154
155
156class ClrBase:
157    """Base class for color convenience class."""
158
159    RST: ClassVar[str]
160    BLD: ClassVar[str]
161    UND: ClassVar[str]
162    INV: ClassVar[str]
163
164    # Normal foreground colors
165    BLK: ClassVar[str]
166    RED: ClassVar[str]
167    GRN: ClassVar[str]
168    YLW: ClassVar[str]
169    BLU: ClassVar[str]
170    MAG: ClassVar[str]
171    CYN: ClassVar[str]
172    WHT: ClassVar[str]
173
174    # Normal background colors.
175    BBLK: ClassVar[str]
176    BRED: ClassVar[str]
177    BGRN: ClassVar[str]
178    BYLW: ClassVar[str]
179    BBLU: ClassVar[str]
180    BMAG: ClassVar[str]
181    BCYN: ClassVar[str]
182    BWHT: ClassVar[str]
183
184    # Strong foreground colors
185    SBLK: ClassVar[str]
186    SRED: ClassVar[str]
187    SGRN: ClassVar[str]
188    SYLW: ClassVar[str]
189    SBLU: ClassVar[str]
190    SMAG: ClassVar[str]
191    SCYN: ClassVar[str]
192    SWHT: ClassVar[str]
193
194    # Strong background colors.
195    SBBLK: ClassVar[str]
196    SBRED: ClassVar[str]
197    SBGRN: ClassVar[str]
198    SBYLW: ClassVar[str]
199    SBBLU: ClassVar[str]
200    SBMAG: ClassVar[str]
201    SBCYN: ClassVar[str]
202    SBWHT: ClassVar[str]
203
204
205class ClrAlways(ClrBase):
206    """Convenience class for color terminal output.
207
208    This version has colors always enabled. Generally you should use Clr which
209    points to the correct enabled/disabled class depending on the environment.
210    """
211
212    color_enabled = True
213
214    # Styles
215    RST = TerminalColor.RESET.value
216    BLD = TerminalColor.BOLD.value
217    UND = TerminalColor.UNDERLINE.value
218    INV = TerminalColor.INVERSE.value
219
220    # Normal foreground colors
221    BLK = TerminalColor.BLACK.value
222    RED = TerminalColor.RED.value
223    GRN = TerminalColor.GREEN.value
224    YLW = TerminalColor.YELLOW.value
225    BLU = TerminalColor.BLUE.value
226    MAG = TerminalColor.MAGENTA.value
227    CYN = TerminalColor.CYAN.value
228    WHT = TerminalColor.WHITE.value
229
230    # Normal background colors.
231    BBLK = TerminalColor.BG_BLACK.value
232    BRED = TerminalColor.BG_RED.value
233    BGRN = TerminalColor.BG_GREEN.value
234    BYLW = TerminalColor.BG_YELLOW.value
235    BBLU = TerminalColor.BG_BLUE.value
236    BMAG = TerminalColor.BG_MAGENTA.value
237    BCYN = TerminalColor.BG_CYAN.value
238    BWHT = TerminalColor.BG_WHITE.value
239
240    # Strong foreground colors
241    SBLK = TerminalColor.STRONG_BLACK.value
242    SRED = TerminalColor.STRONG_RED.value
243    SGRN = TerminalColor.STRONG_GREEN.value
244    SYLW = TerminalColor.STRONG_YELLOW.value
245    SBLU = TerminalColor.STRONG_BLUE.value
246    SMAG = TerminalColor.STRONG_MAGENTA.value
247    SCYN = TerminalColor.STRONG_CYAN.value
248    SWHT = TerminalColor.STRONG_WHITE.value
249
250    # Strong background colors.
251    SBBLK = TerminalColor.STRONG_BG_BLACK.value
252    SBRED = TerminalColor.STRONG_BG_RED.value
253    SBGRN = TerminalColor.STRONG_BG_GREEN.value
254    SBYLW = TerminalColor.STRONG_BG_YELLOW.value
255    SBBLU = TerminalColor.STRONG_BG_BLUE.value
256    SBMAG = TerminalColor.STRONG_BG_MAGENTA.value
257    SBCYN = TerminalColor.STRONG_BG_CYAN.value
258    SBWHT = TerminalColor.STRONG_BG_WHITE.value
259
260
261class ClrNever(ClrBase):
262    """Convenience class for color terminal output.
263
264    This version has colors disabled. Generally you should use Clr which
265    points to the correct enabled/disabled class depending on the environment.
266    """
267
268    color_enabled = False
269
270    # Styles
271    RST = ''
272    BLD = ''
273    UND = ''
274    INV = ''
275
276    # Normal foreground colors
277    BLK = ''
278    RED = ''
279    GRN = ''
280    YLW = ''
281    BLU = ''
282    MAG = ''
283    CYN = ''
284    WHT = ''
285
286    # Normal background colors.
287    BBLK = ''
288    BRED = ''
289    BGRN = ''
290    BYLW = ''
291    BBLU = ''
292    BMAG = ''
293    BCYN = ''
294    BWHT = ''
295
296    # Strong foreground colors
297    SBLK = ''
298    SRED = ''
299    SGRN = ''
300    SYLW = ''
301    SBLU = ''
302    SMAG = ''
303    SCYN = ''
304    SWHT = ''
305
306    # Strong background colors.
307    SBBLK = ''
308    SBRED = ''
309    SBGRN = ''
310    SBYLW = ''
311    SBBLU = ''
312    SBMAG = ''
313    SBCYN = ''
314    SBWHT = ''
315
316
317_envval = os.environ.get('EFRO_TERMCOLORS')
318color_enabled: bool = (
319    True
320    if _envval == '1'
321    else False if _envval == '0' else _default_color_enabled()
322)
323Clr: type[ClrBase] = ClrAlways if color_enabled else ClrNever