• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright (c) 2012 Giorgos Verigakis <verigak@gmail.com>
2#
3# Permission to use, copy, modify, and distribute this software for any
4# purpose with or without fee is hereby granted, provided that the above
5# copyright notice and this permission notice appear in all copies.
6#
7# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
8# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
9# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
10# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
11# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
12# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
13# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
14
15from functools import partial
16from typing import List, Optional, Union
17
18
19# ANSI color names. There is also a "default"
20COLORS = ('black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white')
21
22# ANSI style names
23STYLES = (
24    'none',
25    'bold',
26    'faint',
27    'italic',
28    'underline',
29    'blink',
30    'blink2',
31    'negative',
32    'concealed',
33    'crossed',
34)
35
36
37ColorSpec = Union[str, int]
38
39
40def _join(*values: ColorSpec) -> str:
41    return ';'.join(str(v) for v in values)
42
43
44def _color_code(spec: ColorSpec, base: int) -> str:
45    if isinstance(spec, str):
46        spec = spec.strip().lower()
47
48    if spec == 'default':
49        return _join(base + 9)
50    elif spec in COLORS:
51        return _join(base + COLORS.index(spec))
52    elif isinstance(spec, int) and 0 <= spec <= 255:
53        return _join(base + 8, 5, spec)
54    else:
55        raise ValueError('Invalid color spec "%s"' % spec)
56
57
58def color(
59    s: str,
60    fg: Optional[ColorSpec] = None,
61    bg: Optional[ColorSpec] = None,
62    style: Optional[str] = None,
63) -> str:
64    codes: List[ColorSpec] = []
65
66    if fg:
67        codes.append(_color_code(fg, 30))
68    if bg:
69        codes.append(_color_code(bg, 40))
70    if style:
71        for style_part in style.split('+'):
72            if style_part in STYLES:
73                codes.append(STYLES.index(style_part))
74            else:
75                raise ValueError('Invalid style "%s"' % style_part)
76
77    if codes:
78        return '\x1b[{0}m{1}\x1b[0m'.format(_join(*codes), s)
79    else:
80        return s
81
82
83# Foreground color shortcuts
84black = partial(color, fg='black')
85red = partial(color, fg='red')
86green = partial(color, fg='green')
87yellow = partial(color, fg='yellow')
88blue = partial(color, fg='blue')
89magenta = partial(color, fg='magenta')
90cyan = partial(color, fg='cyan')
91white = partial(color, fg='white')
92
93# Style shortcuts
94bold = partial(color, style='bold')
95none = partial(color, style='none')
96faint = partial(color, style='faint')
97italic = partial(color, style='italic')
98underline = partial(color, style='underline')
99blink = partial(color, style='blink')
100blink2 = partial(color, style='blink2')
101negative = partial(color, style='negative')
102concealed = partial(color, style='concealed')
103crossed = partial(color, style='crossed')
104