1# -*- coding: utf-8 -*- 2# Copyright (c) 2011 The Chromium OS Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6"""Terminal utilities 7 8This module handles terminal interaction including ANSI color codes. 9""" 10 11from __future__ import print_function 12 13import os 14import sys 15 16from autotest_lib.utils.frozen_chromite.lib import cros_build_lib 17 18 19class Color(object): 20 """Conditionally wraps text in ANSI color escape sequences.""" 21 BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8) 22 BOLD = -1 23 COLOR_START = '\033[1;%dm' 24 BOLD_START = '\033[1m' 25 RESET = '\033[0m' 26 27 def __init__(self, enabled=None): 28 """Create a new Color object, optionally disabling color output. 29 30 Args: 31 enabled: True if color output should be enabled. If False then this 32 class will not add color codes at all. 33 """ 34 self._enabled = enabled 35 if self._enabled is None: 36 self._enabled = self.UserEnabled() 37 if self._enabled is None: 38 self._enabled = sys.stdout.isatty() 39 40 def Start(self, color): 41 """Returns a start color code. 42 43 Args: 44 color: Color to use, .e.g BLACK, RED, etc. 45 46 Returns: 47 If color is enabled, returns an ANSI sequence to start the given color, 48 otherwise returns empty string 49 """ 50 if self._enabled: 51 return self.COLOR_START % (color + 30) 52 return '' 53 54 def Stop(self): 55 """Returns a stop color code. 56 57 Returns: 58 If color is enabled, returns an ANSI color reset sequence, otherwise 59 returns empty string 60 """ 61 if self._enabled: 62 return self.RESET 63 return '' 64 65 def Color(self, color, text): 66 """Returns text with conditionally added color escape sequences. 67 68 Keyword arguments: 69 color: Text color -- one of the color constants defined in this class. 70 text: The text to color. 71 72 Returns: 73 If self._enabled is False, returns the original text. If it's True, 74 returns text with color escape sequences based on the value of color. 75 """ 76 if not self._enabled: 77 return text 78 if color == self.BOLD: 79 start = self.BOLD_START 80 else: 81 start = self.COLOR_START % (color + 30) 82 return start + text + self.RESET 83 84 @staticmethod 85 def UserEnabled(): 86 """See if the global colorization preference is enabled ($NOCOLOR env)""" 87 is_disabled = cros_build_lib.BooleanShellValue( 88 os.environ.get('NOCOLOR'), msg='$NOCOLOR env var is invalid', 89 default=None) 90 return not is_disabled if is_disabled is not None else None 91