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