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
8 # http://www.apache.org/licenses/LICENSE-2.0
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.
17 Terminal colorization utilities.
20 from StringIO import StringIO
26 from ..utils.formatting import safe_str
29 def _restore_terminal():
34 atexit.register(_restore_terminal)
37 class StringStylizer(object):
38 def __init__(self, str_, color_spec=None):
40 self._color_spec = color_spec
44 return '{schema}{str}{reset}'.format(
45 schema=self._color_spec, str=safe_str(self._str), reset=Colors.Style.RESET_ALL)
48 def __add__(self, other):
49 return safe_str(self) + other
51 def __radd__(self, other):
52 return other + safe_str(self)
54 def color(self, color_spec):
55 self._color_spec = color_spec
57 def replace(self, old, new, **kwargs):
58 self._str = self._str.replace(safe_str(old), safe_str(new), **kwargs)
60 def format(self, *args, **kwargs):
61 self._str = self._str.format(*args, **kwargs)
63 def highlight(self, pattern, schema):
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)
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))
79 Style = colorama.Style
82 'fore': dict(_get_colors(Fore)),
83 'back': dict(_get_colors(Back)),
84 'style': dict(_get_colors(Style))
88 class ColorSpec(object):
89 def __init__(self, fore=None, back=None, style=None):
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.
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)
102 return self._str.getvalue()
104 def __add__(self, other):
105 return str(self) + str(other)
107 def __radd__(self, other):
108 return str(other) + str(self)