slither.utils.colors

 1from functools import partial
 2import platform
 3import sys
 4
 5
 6class Colors:  # pylint: disable=too-few-public-methods
 7    COLORIZATION_ENABLED = True
 8    RED = "\033[91m"
 9    GREEN = "\033[92m"
10    YELLOW = "\033[93m"
11    BLUE = "\033[94m"
12    MAGENTA = "\033[95m"
13    BOLD = "\033[1m"
14    END = "\033[0m"
15
16
17def colorize(color: Colors, txt: str) -> str:
18    if Colors.COLORIZATION_ENABLED:
19        return f"{color}{txt}{Colors.END}"
20    return txt
21
22
23def enable_windows_virtual_terminal_sequences() -> bool:
24    """
25    Sets the appropriate flags to enable virtual terminal sequences in a Windows command prompt.
26    Reference: https://docs.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences
27    """
28
29    try:
30        # pylint: disable=import-outside-toplevel
31        from ctypes import windll, byref  # type: ignore
32        from ctypes.wintypes import DWORD, HANDLE
33
34        kernel32 = windll.kernel32
35        virtual_terminal_flag = 0x04  # ENABLE_VIRTUAL_TERMINAL_PROCESSING
36
37        # Obtain our stdout/stderr handles.
38        # Reference: https://docs.microsoft.com/en-us/windows/console/getstdhandle
39        handle_stdout = kernel32.GetStdHandle(-11)
40        handle_stderr = kernel32.GetStdHandle(-12)
41
42        # Loop for each stdout/stderr handle.
43        for current_handle in [handle_stdout, handle_stderr]:
44
45            # If we get a null handle, or fail any subsequent calls in this scope, we do not colorize any output.
46            if current_handle is None or current_handle == HANDLE(-1):
47                return False
48
49            # Try to obtain the current flags for the console.
50            current_mode = DWORD()
51            if not kernel32.GetConsoleMode(current_handle, byref(current_mode)):
52                return False
53
54            # If the virtual terminal sequence processing is not yet enabled, we enable it.
55            if (current_mode.value & virtual_terminal_flag) == 0:
56                if not kernel32.SetConsoleMode(
57                    current_handle, current_mode.value | virtual_terminal_flag
58                ):
59                    return False
60    except Exception:  # pylint: disable=broad-except
61        # Any generic failure (possibly from calling these methods on older Windows builds where they do not exist)
62        # will fall back onto disabling colorization.
63        return False
64
65    return True
66
67
68def set_colorization_enabled(enabled: bool) -> None:
69    """
70    Sets the enabled state of output colorization.
71    :param enabled: Boolean indicating whether output should be colorized.
72    :return: None
73    """
74    # If color is supposed to be enabled and this is windows, we have to enable console virtual terminal sequences:
75    if enabled and platform.system() == "Windows":
76        Colors.COLORIZATION_ENABLED = enable_windows_virtual_terminal_sequences()
77    else:
78        # This is not windows, or colorization is being disabled, so we can adjust the state immediately.
79        Colors.COLORIZATION_ENABLED = enabled
80
81
82green = partial(colorize, Colors.GREEN)
83yellow = partial(colorize, Colors.YELLOW)
84red = partial(colorize, Colors.RED)
85blue = partial(colorize, Colors.BLUE)
86magenta = partial(colorize, Colors.MAGENTA)
87bold = partial(colorize, Colors.BOLD)
88
89# We enable colorization by default if the output is a tty
90set_colorization_enabled(sys.stdout.isatty())