PEP8: add spaces after operators
[nivanova/samba-autobuild/.git] / python / samba / colour.py
1 # ANSI codes for 4 bit and xterm-256color
2 #
3 # Copyright (C) Andrew Bartlett 2018
4 #
5 # Originally written by Douglas Bagnall
6 #
7 # This program is free software; you can redistribute it and/or modify
8 # it under the terms of the GNU General Public License as published by
9 # the Free Software Foundation; either version 3 of the License, or
10 # (at your option) any later version.
11 #
12 # This program is distributed in the hope that it will be useful,
13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15 # GNU General Public License for more details.
16 #
17 # You should have received a copy of the GNU General Public License
18 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
19 #
20 # The 4 bit colours are available as global variables with names like
21 # RED, DARK_RED, REV_RED (for red background), and REV_DARK_RED. If
22 # switch_colour_off() is called, these names will all point to the
23 # empty string. switch_colour_on() restores the default values.
24 #
25 # The 256-colour codes are obtained using xterm_256_color(n), where n
26 # is the number of the desired colour.
27
28
29 def _gen_ansi_colours():
30     g = globals()
31     for i, name in enumerate(('BLACK', 'RED', 'GREEN', 'YELLOW', 'BLUE',
32                               'MAGENTA', 'CYAN', 'WHITE')):
33         g[name] = "\033[1;3%dm" % i
34         g['DARK_' + name] = "\033[3%dm" % i
35         g['REV_' + name] = "\033[1;4%dm" % i
36         g['REV_DARK_' + name] = "\033[4%dm" % i
37
38     # kcc.debug uses these aliases (which make visual sense)
39     g['PURPLE'] = DARK_MAGENTA
40     g['GREY'] = DARK_WHITE
41
42     # C_NORMAL resets to normal, whatever that is
43     g['C_NORMAL'] = "\033[0m"
44
45     # Non-colour ANSI codes.
46     g['UNDERLINE'] = "\033[4m"
47
48
49 _gen_ansi_colours()
50
51 # Generate functions that colour a string. The functions look like
52 # this:
53 #
54 #    c_BLUE("hello")  # "\033[1;34mhello\033[0m" -> blue text
55 #    c_DARK_RED(3)    # 3 will be stringified and coloured
56 #
57 # but if colour is switched off, no colour codes are added.
58 #
59 #    c_BLUE("hello")  # "hello"
60 #
61 # The definition of the functions looks a little odd, because we want
62 # to bake in the name of the colour but not its actual value.
63
64 for _k in list(globals().keys()):
65     if _k.isupper():
66         def _f(s, name=_k):
67             return "%s%s%s" % (globals()[name], s, C_NORMAL)
68         globals()['c_%s' % _k] = _f
69
70 del _k, _f
71
72
73 def switch_colour_off():
74     """Convert all the ANSI colour codes into empty strings."""
75     g = globals()
76     for k, v in list(g.items()):
77         if k.isupper() and isinstance(v, str) and v.startswith('\033'):
78             g[k] = ''
79
80
81 def switch_colour_on():
82     """Regenerate all the ANSI colour codes."""
83     _gen_ansi_colours()
84
85
86 def xterm_256_colour(n, bg=False, bold=False):
87     weight = '01;' if bold else ''
88     target = '48' if bg else '38'
89
90     return "\033[%s%s;5;%dm" % (weight, target, int(n))