# -*- coding: utf-8 -*- """A simple and elegant wrapper for colorama.""" import os from random import choice, seed import re import sys import colorama PY3 = sys.version_info[0] >= 3 COLORS = ( 'red', 'green', 'yellow', 'blue', 'black', 'magenta', 'cyan', 'white' ) __all__ = COLORS + ( 'normal', 'clean', 'disable', 'enable', 'random', 'replace_colors', 'reset_replace_colors' ) colorama.init() REPLACE_COLORS = {} seed() if 'get_ipython' in dir(): """ when ipython is fired lot of variables like _oh, etc are used. There are so many ways to find current python interpreter is ipython. get_ipython is easiest is most appealing for readers to understand. """ DISABLE_COLOR = True else: DISABLE_COLOR = False if os.getenv("TERM") == "dumb": DISABLE_COLOR = True class ColoredString(object): """Enhanced string for __len__ operations on Colored output.""" def __init__(self, color, s, always_color=False, bold=False): super(ColoredString, self).__init__() if not PY3 and isinstance(s, unicode): # noqa: F821 self.s = s.encode('utf-8') else: self.s = str(s) self.color = REPLACE_COLORS.get(color, color) self.always_color = always_color self.bold = bold if os.environ.get('CLINT_FORCE_COLOR'): self.always_color = True def __getattr__(self, att): def func_help(*args, **kwargs): result = getattr(self.s, att)(*args, **kwargs) try: is_result_string = isinstance(result, basestring) except NameError: is_result_string = isinstance(result, str) if is_result_string: return self._new(result) elif isinstance(result, list): return [self._new(x) for x in result] else: return result return func_help @property def color_str(self): style = 'BRIGHT' if self.bold else 'NORMAL' c = '%s%s%s%s%s' % (getattr(colorama.Fore, self.color), getattr(colorama.Style, style), self.s, colorama.Fore.RESET, colorama.Style.NORMAL) if self.always_color: return c elif hasattr(sys.stdout, 'isatty') \ and sys.stdout.isatty() \ and not DISABLE_COLOR: return c else: return self.s def __len__(self): return len(self.s) def __repr__(self): return "<%s-string: '%s'>" % (self.color, self.s) def __unicode__(self): value = self.color_str if isinstance(value, bytes): return value.decode('utf8') return value if PY3: __str__ = __unicode__ else: def __str__(self): return self.color_str def __iter__(self): return iter(self.color_str) def __add__(self, other): return str(self.color_str) + str(other) def __radd__(self, other): return str(other) + str(self.color_str) def __mul__(self, other): return (self.color_str * other) def _new(self, s): return ColoredString(self.color, s) _colors = {x: x.upper() for x in COLORS} _colors['normal'] = 'RESET' for key, val in _colors.items(): function = eval( 'lambda s, always=False, bold=False: ColoredString("{}", s, always_color=always, bold=bold)'.format(val)) # noqa: E501 locals()[key] = function del key, val, _colors, function def clean(s): strip = re.compile(r'(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]') txt = strip.sub('', s) return txt def random(string, always=False, bold=False, colors=COLORS): """Selects a color at random from a list.""" colors = list(filter(lambda color: color in COLORS, colors)) or COLORS return ColoredString( choice(colors).upper(), string, always_color=always, bold=bold ) def disable(): """Disables colors.""" global DISABLE_COLOR DISABLE_COLOR = True def enable(): """Enables colors.""" global DISABLE_COLOR DISABLE_COLOR = False def replace_colors(replace_dict): """Replace colors to customize the look under specific background.""" global REPLACE_COLORS assert isinstance(replace_dict, dict) REPLACE_COLORS = {k.upper(): v.upper() for k, v in replace_dict.items()} def reset_replace_colors(): """Reset any set replace colors.""" global REPLACE_COLORS REPLACE_COLORS = {}