| # Copyright 2011 The ChromiumOS Authors |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Terminal utilities |
| |
| This module handles terminal interaction including ANSI color codes. |
| """ |
| |
| import os |
| import sys |
| from typing import Optional |
| |
| from chromite.lib import cros_build_lib |
| |
| |
| class Color: |
| """Conditionally wraps text in ANSI color escape sequences.""" |
| |
| BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8) |
| BOLD = -1 |
| COLOR_START = "\033[1;%dm" |
| BACKGROUND_START = "\u001b[%dm" |
| BOLD_START = "\033[1m" |
| RESET = "\033[0m" |
| BACKGROUND_RESET = "\u001b[0m" |
| |
| def __init__(self, enabled: Optional[bool] = None) -> None: |
| """Create a new Color object, optionally disabling color output. |
| |
| Args: |
| enabled: True if color output should be enabled. If False then this |
| class will not add color codes at all. |
| """ |
| self._enabled = enabled |
| if self._enabled is None: |
| self._enabled = self.UserEnabled() |
| if self._enabled is None: |
| self._enabled = sys.stdout.isatty() |
| |
| def Start(self, color: int) -> str: |
| """Returns a start color code. |
| |
| Args: |
| color: Color to use, .e.g BLACK, RED, etc. |
| |
| Returns: |
| If color is enabled, returns an ANSI sequence to start the given |
| color, otherwise returns empty string |
| """ |
| if self._enabled: |
| return self.COLOR_START % (color + 30) |
| return "" |
| |
| def Stop(self) -> str: |
| """Returns a stop color code. |
| |
| Returns: |
| If color is enabled, returns an ANSI color reset sequence, otherwise |
| returns empty string |
| """ |
| if self._enabled: |
| return self.RESET |
| return "" |
| |
| def Color( |
| self, color: int, text: str, background_color: Optional[int] = None |
| ) -> str: |
| """Returns text with conditionally added color escape sequences. |
| |
| Keyword arguments: |
| color: Text color -- one of the color constants defined in this |
| class. |
| text: The text to color. |
| background_color: Background highlight color -- one of the color |
| constants defined in this class. |
| |
| Returns: |
| If self._enabled is False, returns the original text. If it's True, |
| returns text with color escape sequences based on the value of |
| color. |
| """ |
| if not self._enabled: |
| return text |
| if color == self.BOLD: |
| start = self.BOLD_START |
| else: |
| start = self.COLOR_START % (color + 30) |
| end = self.RESET |
| if background_color: |
| start += self.BACKGROUND_START % (color + 30) |
| end += self.RESET |
| return start + text + end |
| |
| @staticmethod |
| def UserEnabled() -> Optional[bool]: |
| """See if the global colorization preference is enabled. |
| |
| Uses the $NOCOLOR envvar. |
| """ |
| is_disabled = cros_build_lib.BooleanShellValue( |
| os.environ.get("NOCOLOR"), |
| msg="$NOCOLOR env var is invalid", |
| default=None, |
| ) |
| return not is_disabled if is_disabled is not None else None |