1# Copyright (c) 2012 Giorgos Verigakis <[email protected]>
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
19class ColorError(ValueError):
20    """Error raised when a color spec is invalid."""
21
22
23# ANSI color names. There is also a "default"
24COLORS = ('black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white')
25
26# ANSI style names
27STYLES = (
28    'none',
29    'bold',
30    'faint',
31    'italic',
32    'underline',
33    'blink',
34    'blink2',
35    'negative',
36    'concealed',
37    'crossed',
38)
39
40
41ColorSpec = Union[str, int]
42
43
44def _join(*values: ColorSpec) -> str:
45    return ';'.join(str(v) for v in values)
46
47
48def _color_code(spec: ColorSpec, base: int) -> str:
49    if isinstance(spec, str):
50        spec = spec.strip().lower()
51
52    if spec == 'default':
53        return _join(base + 9)
54    elif spec in COLORS:
55        return _join(base + COLORS.index(spec))
56    elif isinstance(spec, int) and 0 <= spec <= 255:
57        return _join(base + 8, 5, spec)
58    else:
59        raise ColorError('Invalid color spec "%s"' % spec)
60
61
62def color(
63    s: str,
64    fg: Optional[ColorSpec] = None,
65    bg: Optional[ColorSpec] = None,
66    style: Optional[str] = None,
67) -> str:
68    codes: List[ColorSpec] = []
69
70    if fg:
71        codes.append(_color_code(fg, 30))
72    if bg:
73        codes.append(_color_code(bg, 40))
74    if style:
75        for style_part in style.split('+'):
76            if style_part in STYLES:
77                codes.append(STYLES.index(style_part))
78            else:
79                raise ColorError('Invalid style "%s"' % style_part)
80
81    if codes:
82        return '\x1b[{0}m{1}\x1b[0m'.format(_join(*codes), s)
83    else:
84        return s
85
86
87# Foreground color shortcuts
88black = partial(color, fg='black')
89red = partial(color, fg='red')
90green = partial(color, fg='green')
91yellow = partial(color, fg='yellow')
92blue = partial(color, fg='blue')
93magenta = partial(color, fg='magenta')
94cyan = partial(color, fg='cyan')
95white = partial(color, fg='white')
96
97# Style shortcuts
98bold = partial(color, style='bold')
99none = partial(color, style='none')
100faint = partial(color, style='faint')
101italic = partial(color, style='italic')
102underline = partial(color, style='underline')
103blink = partial(color, style='blink')
104blink2 = partial(color, style='blink2')
105negative = partial(color, style='negative')
106concealed = partial(color, style='concealed')
107crossed = partial(color, style='crossed')
108