azure/aria/aria-extension-cloudify/src/aria/aria/cli/color.py

   1 # Licensed to the Apache Software Foundation (ASF) under one or more
   2 # contributor license agreements.  See the NOTICE file distributed with
   3 # this work for additional information regarding copyright ownership.
   4 # The ASF licenses this file to You under the Apache License, Version 2.0
   5 # (the "License"); you may not use this file except in compliance with
   6 # the License.  You may obtain a copy of the License at
   7 #
   8 #     http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 """
  17 Terminal colorization utilities.
  18 """
  19
  20 from StringIO import StringIO
  21 import atexit
  22 import re
  23
  24 import colorama
  25
  26 from ..utils.formatting import safe_str
  27
  28
  29 def _restore_terminal():
  30     colorama.deinit()
  31
  32
  33 colorama.init()
  34 atexit.register(_restore_terminal)
  35
  36
  37 class StringStylizer(object):
  38     def __init__(self, str_, color_spec=None):
  39         self._str = str_
  40         self._color_spec = color_spec
  41
  42     def __repr__(self):
  43         if self._color_spec:
  44             return '{schema}{str}{reset}'.format(
  45                 schema=self._color_spec, str=safe_str(self._str), reset=Colors.Style.RESET_ALL)
  46         return self._str
  47
  48     def __add__(self, other):
  49         return safe_str(self) + other
  50
  51     def __radd__(self, other):
  52         return other + safe_str(self)
  53
  54     def color(self, color_spec):
  55         self._color_spec = color_spec
  56
  57     def replace(self, old, new, **kwargs):
  58         self._str = self._str.replace(safe_str(old), safe_str(new), **kwargs)
  59
  60     def format(self, *args, **kwargs):
  61         self._str = self._str.format(*args, **kwargs)
  62
  63     def highlight(self, pattern, schema):
  64         if pattern is None:
  65             return
  66         for match in set(re.findall(re.compile(pattern), self._str)):
  67             self.replace(match, schema + match + Colors.Style.RESET_ALL + self._color_spec)
  68
  69
  70 def _get_colors(color_type):
  71     for name in dir(color_type):
  72         if not name.startswith('_'):
  73             yield (name.lower(), getattr(color_type, name))
  74
  75
  76 class Colors(object):
  77     Fore = colorama.Fore
  78     Back = colorama.Back
  79     Style = colorama.Style
  80
  81     _colors = {
  82         'fore': dict(_get_colors(Fore)),
  83         'back': dict(_get_colors(Back)),
  84         'style': dict(_get_colors(Style))
  85     }
  86
  87
  88 class ColorSpec(object):
  89     def __init__(self, fore=None, back=None, style=None):
  90         """
  91         It is possible to provide fore, back and style arguments. Each could be either the color as
  92         lowercase letters, or the full color name for Colorama.
  93         """
  94         self._kwargs = dict(fore=fore, back=back, style=style)
  95         self._str = StringIO()
  96         for type_, colors in Colors._colors.items():
  97             value = self._kwargs.get(type_, None)
  98             # the former case is if the value is a string, the latter is in case of an object.
  99             self._str.write(colors.get(value) or value)
 100
 101     def __str__(self):
 102         return self._str.getvalue()
 103
 104     def __add__(self, other):
 105         return str(self) + str(other)
 106
 107     def __radd__(self, other):
 108         return str(other) + str(self)