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))
  91
  92
  93 def is_colour_wanted(*streams, hint='auto'):
  94     """The hint is presumably a --color argument.
  95
  96     The streams to be considered can be file objects or file names,
  97     with '-' being a special filename indicating stdout.
  98
  99     We follow the behaviour of GNU `ls` in what we accept.
 100     * `git` is stricter, accepting only {always,never,auto}.
 101     * `grep` is looser, accepting mixed case variants.
 102     * historically we have used {yes,no,auto}.
 103     * {always,never,auto} appears the commonest convention.
 104     * if the caller tries to opt out of choosing and sets hint to None
 105       or '', we assume 'auto'.
 106     """
 107     if hint in ('no', 'never', 'none'):
 108         return False
 109
 110     if hint in ('yes', 'always', 'force'):
 111         return True
 112
 113     if hint not in ('auto', 'tty', 'if-tty', None, ''):
 114         raise ValueError("unexpected colour hint: {hint}; "
 115                          "try always|never|auto")
 116
 117     from os import environ
 118     if environ.get('NO_COLOR'):
 119         # Note: per spec, we treat the empty string as if unset.
 120         return False
 121
 122     for stream in streams:
 123         if isinstance(stream, str):
 124             # This function can be passed filenames instead of file
 125             # objects, in which case we treat '-' as stdout, and test
 126             # that. Any other string is not regarded as a tty.
 127             if stream != '-':
 128                 return False
 129             import sys
 130             stream = sys.stdout
 131
 132         if not stream.isatty():
 133             return False
 134     return True
 135
 136
 137 def colour_if_wanted(*streams, hint='auto'):
 138     wanted = is_colour_wanted(*streams, hint=hint)
 139     if wanted:
 140         switch_colour_on()
 141     else:
 142         switch_colour_off()
 143     return wanted