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(f"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
 144
 145
 146 def colourdiff(a, b):
 147     """Generate a string comparing two strings or byte sequences, with
 148     differences coloured to indicate what changed.
 149
 150     Byte sequences are printed as hex pairs separated by colons.
 151     """
 152     from difflib import SequenceMatcher
 153     out = []
 154     if isinstance(a, bytes):
 155         a = a.hex(':')
 156     if isinstance(b, bytes):
 157         b = b.hex(':')
 158     a = a.replace(' ', '␠')
 159     b = b.replace(' ', '␠')
 160
 161     s = SequenceMatcher(None, a, b)
 162     for op, al, ar, bl, br in s.get_opcodes():
 163         if op == 'equal':
 164             out.append(a[al: ar])
 165         elif op == 'delete':
 166             out.append(c_RED(a[al: ar]))
 167         elif op == 'insert':
 168             out.append(c_GREEN(b[bl: br]))
 169         elif op == 'replace':
 170             out.append(c_RED(a[al: ar]))
 171             out.append(c_GREEN(b[bl: br]))
 172         else:
 173             out.append(f' --unknown diff op {op}!-- ')
 174
 175     return ''.join(out)