File indexing completed on 2024-11-03 08:24:24

0001 # -*- coding: UTF-8 -*-
0002 """
0003 Standard codes for terminal colors.
0004 
0005 @author: Chusslove Illich <caslav.ilic@gmx.net>
0006 @author: Sébastien Renard <sebastien.renard@digitalfox.org>
0007 @license: GPLv3
0008 """
0009 
0010 from optparse import OptionParser
0011 import re
0012 import sys
0013 
0014 # NOTE: Must not import anything from Pology, as top __init__ includes this.
0015 
0016 
0017 _xml_entities = {
0018     "lt": "<",
0019     "gt": ">",
0020     "apos": "'",
0021     "quot": "\"",
0022     "amp": "&",
0023 }
0024 
0025 def _resolve_xml_ents (text):
0026 
0027     segs = []
0028     p = 0
0029     while True:
0030         p1 = p
0031         p = text.find("&", p1)
0032         if p < 0:
0033             segs.append(text[p1:])
0034             break
0035         segs.append(text[p1:p])
0036         p2 = p
0037         p = text.find(";", p2)
0038         if p < 0:
0039             segs.append(text[p2:])
0040             break
0041         ent = text[p2 + 1:p]
0042         val = _xml_entities.get(ent)
0043         if val is None:
0044             segs.append(text[p2])
0045             p = p2 + 1
0046         else:
0047             segs.append(val)
0048             p += 1
0049     rtext = "".join(segs)
0050     return rtext
0051 
0052 
0053 def _escape_xml_ents (text):
0054 
0055     rtext = text.replace("&", "&amp;")
0056     for ent, val in list(_xml_entities.items()):
0057         if val != "&":
0058             rtext = rtext.replace(val, "&" + ent + ";")
0059     return rtext
0060 
0061 
0062 class ColorString (str):
0063     """
0064     Class for strings with color markup.
0065 
0066     This class provides automatic resolution of color XML markup
0067     in strings for various output formats.
0068     It automatically escapes any raw strings combined with it
0069     (e.g. when using the C{%} or C{+} operators)
0070     and returns objects of its own type from methods
0071     (e.g. from C{split()} or C{strip()}).
0072     Otherwise it should behave like a normal string.
0073 
0074     Note that usage of this class is expensive,
0075     given that arguments are constantly checked and strings escaped.
0076     It should be used only for user-visible output,
0077     i.e. where human reaction time is the limiting factor.
0078     """
0079 
0080     def _escape (self, v):
0081         if isinstance(v, str) and not isinstance(v, ColorString):
0082             v = str(v)
0083             v = _escape_xml_ents(v)
0084 
0085         return v
0086 
0087 
0088     def __add__ (self, other):
0089 
0090         return ColorString(str.__add__(self, self._escape(other)))
0091 
0092 
0093     def __radd__ (self, other):
0094 
0095         return ColorString(str.__add__(self._escape(other), self))
0096 
0097 
0098     def __mod__ (self, args):
0099 
0100         if isinstance(args, dict):
0101             rargs = dict((k, self._escape(v)) for k, v in list(args.items()))
0102         elif isinstance(args, tuple):
0103             rargs = tuple(self._escape(v) for v in args)
0104         else:
0105             rargs = self._escape(args)
0106         return ColorString(str.__mod__(self, rargs))
0107 
0108 
0109     def __repr__ (self):
0110 
0111         return "%s(%s)" % (self.__class__.__name__, str.__repr__(self))
0112 
0113     def __iter__(self):
0114         for c in str(self):
0115             yield self.__class__(c)
0116 
0117 
0118     def join (self, strings):
0119         rstrings = [self._escape(s) for s in strings]
0120         return ColorString(str.join(self, rstrings))
0121 
0122 
0123     def resolve (self, ctype=None, dest=None):
0124         """
0125         Resolve color markup according to given type and destination.
0126 
0127         Currently available coloring types (values of C{ctype} parameter):
0128           - C{"none"}: no coloring
0129           - C{"term"}: ANSI color escape sequences (for terminal output)
0130           - C{"html"}: HTML markup (for integration into HTML pages)
0131         If C{ctype} is C{None}, it is taken from global coloring options.
0132 
0133         Some coloring types may be applied conditionally, based on whether
0134         the intended output destination is a file or terminal.
0135         If this is desired, the file descriptor of the destination
0136         can be given by the C{dest} parameter.
0137 
0138         @param ctype: type of coloring
0139         @type ctype: string
0140         @param dest: destination file descriptor
0141         @type dest: file
0142         @returns: plain string with resolved markup
0143         @rtype: string
0144         """
0145 
0146         # Resolve coloring type, considering all things.
0147         if ctype is None:
0148             ctype = _cglobals.ctype
0149         if ctype in (None, "term"):
0150             if not _cglobals.outdep or (dest and dest.isatty()):
0151                 ctype = "term"
0152             else:
0153                 ctype = "none"
0154 
0155         color_pack = _color_packs.get(ctype)
0156         if color_pack is None:
0157             color_pack = _color_packs.get("none")
0158         colorf, escapef, finalf = color_pack
0159         text = str(self)
0160         rtext, epos = self._resolve_markup_w(text, len(text), 0, None, None,
0161                                              colorf, escapef)
0162         rtext = finalf(rtext)
0163         return rtext
0164 
0165 
0166     def _resolve_markup_w (self, text, tlen, pos, tag, ptag, colorf, escapef):
0167 
0168         rsegs = []
0169         p = pos
0170         valid = True
0171         closed = False
0172         while p < tlen:
0173             pp = p
0174             p = text.find("<", p)
0175             if p < 0:
0176                 p = tlen
0177             seg = text[pp:p]
0178             rsegs.append(escapef(_resolve_xml_ents(seg)))
0179             if p == tlen:
0180                 break
0181             pp = p
0182             stag, closed, p = self._parse_tag(text, tlen, p)
0183             if stag is not None:
0184                 if not closed:
0185                     rseg, p = self._resolve_markup_w(text, tlen, p, stag, tag,
0186                                                      colorf, escapef)
0187                     rsegs.append(rseg)
0188                 else:
0189                     if tag != stag:
0190                         # Wrong closed tag, declare this span not valid
0191                         # and reposition at the tag start.
0192                         valid = False
0193                         p = pp
0194                     break
0195             else:
0196                 # Not a proper tag start, just take literal < and go on.
0197                 rsegs.append("<")
0198                 p = pp + 1
0199         if tag and not closed:
0200             valid = False
0201 
0202         rtext = "".join(rsegs)
0203         if tag:
0204             if valid:
0205                 rtext = colorf(tag, rtext, ptag)
0206             else:
0207                 # Not proper span, put back opening tag.
0208                 rtext = "<%s>%s" % (tag, rtext)
0209 
0210         return rtext, p
0211 
0212 
0213     def _parse_tag (self, text, tlen, pos):
0214 
0215         # FIXME: No possibility of attributes at the moment.
0216         if tlen is None:
0217             tlen = len(text)
0218         p = pos
0219         tag = None
0220         closed = False
0221         if p < tlen and text[p] == "<":
0222             p += 1
0223             while p < tlen and text[p].isspace():
0224                 p += 1
0225             if p < tlen:
0226                 if text[p] == "/":
0227                     p += 1
0228                     closed = True
0229                 pp = p
0230                 p = text.find(">", p)
0231                 if p < 0:
0232                     p = tlen
0233                 else:
0234                     tag = text[pp:p].strip()
0235                     p += 1
0236         return tag, closed, p
0237 
0238 
0239     def visual_segment (self, pos):
0240         """
0241         Get visual representation of raw segment starting from position.
0242 
0243         This function checks whether the segment of the string starting
0244         at given position has the raw or another visual value,
0245         accounting for markup.
0246         If the visual and raw segments differ, the visual representation
0247         and length of the raw segment are returned.
0248         Otherwise, empty string and zero length are returned.
0249 
0250         @param pos: position where to check for visual segment
0251         @type pos: int
0252         @returns: visual segment and length of underlying raw segment
0253         @rtype: string, int
0254         """
0255 
0256         vis, rlen = "", 0
0257 
0258         c = self[pos:pos + 1]
0259         if c == "<":
0260             pos2 = self.find(">", pos)
0261             if pos2 > 0:
0262                 vis, rlen = "", pos2 + 1 - pos
0263         elif c == "&":
0264             pos2 = self.find(";", pos)
0265             if pos2 > 0:
0266                 ent = self[pos + 1:pos2]
0267                 val = _xml_entities.get(ent)
0268                 if val is not None:
0269                     vis, rlen = val, pos2 + 1 - pos
0270 
0271         return vis, rlen
0272 
0273 
0274 def _fill_color_string_class ():
0275 
0276     def wrap_return_type (method):
0277         def wmethod (self, *args, **kwargs):
0278             res = method(self, *args, **kwargs)
0279             if isinstance(res, str):
0280                 res = ColorString(res)
0281             elif isinstance(res, (tuple, list)):
0282                 res2 = []
0283                 for el in res:
0284                     if isinstance(el, str):
0285                         el = ColorString(el)
0286                     res2.append(el)
0287                 res = type(res)(res2)
0288             return res
0289         return wmethod
0290 
0291     for attrname in (
0292         "__getitem__", "__mul__", "__rmul__",
0293         "capitalize", "center", "expandtabs", "ljust", "lower", "lstrip",
0294         "replace", "rjust", "rsplit", "rstrip", "split", "strip", "swapcase",
0295         "title", "translate", "upper", "zfill",
0296     ):
0297         method = getattr(str, attrname)
0298         setattr(ColorString, attrname, wrap_return_type(method))
0299 
0300 _fill_color_string_class()
0301 
0302 
0303 def cjoin (strings, joiner=""):
0304     """
0305     Join strings into a L{ColorString} if any of them are L{ColorString},
0306     otherwise into type of joiner.
0307 
0308     @param strings: strings to join
0309     @type strings: sequence of strings
0310     @param joiner: string to be inserted between each two strings
0311     @type joiner: string
0312     @returns: concatenation by joiner of all strings
0313     @rtype: type(joiner)/L{ColorString}
0314     """
0315 
0316     if not isinstance(joiner, ColorString):
0317         for s in strings:
0318             if isinstance(s, ColorString):
0319                 joiner = ColorString(joiner)
0320                 break
0321     return joiner.join(strings)
0322 
0323 
0324 def cinterp (format, *args, **kwargs):
0325     """
0326     Interpolate arguments into the format string, producing L{ColorString}
0327     if any of the arguments is L{ColorString}, otherwise type of format string.
0328 
0329     The format string can use either positional format directives,
0330     in which case positional arguments are supplied after it,
0331     or it can use named format directives,
0332     in which case keyword arguments are supplied after it.
0333     If both positional and keyword arguments are following the format string,
0334     the behavior is undefined.
0335 
0336     @param format: string with formatting directives
0337     @type format: string
0338     @returns: interpolated strings
0339     @rtype: type(format)/L{ColorString}
0340     """
0341 
0342     iargs = args or kwargs
0343     if not isinstance(format, ColorString):
0344         for v in (list(iargs.values()) if isinstance(iargs, dict) else iargs):
0345             if isinstance(v, ColorString):
0346                 format = ColorString(format)
0347                 break
0348     return format % iargs
0349 
0350 
0351 class ColorOptionParser (OptionParser):
0352     """
0353     Lightweight wrapper for C{OptionParser} from standard library C{optparse},
0354     to gracefully handle L{ColorString} arguments supplied to its methods.
0355     """
0356 
0357     def _cv (self, val):
0358 
0359         if isinstance(val, ColorString):
0360             val = val.resolve("term", sys.stdout)
0361         elif isinstance(val, (list, tuple)):
0362             val = list(map(self._cv, val))
0363         elif isinstance(val, dict):
0364             val = dict((k, self._cv(v)) for k, v in list(val.items()))
0365         return val
0366 
0367 
0368     def __init__ (self, *args, **kwargs):
0369 
0370         OptionParser.__init__(self, *self._cv(args), **self._cv(kwargs))
0371 
0372 
0373     def add_option (self, *args, **kwargs):
0374 
0375         OptionParser.add_option(self, *self._cv(args), **self._cv(kwargs))
0376 
0377 
0378     # FIXME: More overrides.
0379 
0380 
0381 def get_coloring_types ():
0382     """
0383     List of keywords of all available coloring types.
0384     """
0385 
0386     return list(_color_packs.keys())
0387 
0388 
0389 def set_coloring_globals (ctype="term", outdep=True):
0390     """
0391     Set global options for coloring.
0392 
0393     L{ColorString.resolve} will use the type of coloring given
0394     by C{ctype} here whenever its own C{ctype} is set to C{None}.
0395 
0396     If C{outdep} is set to C{False}, L{ColorString.resolve} will not
0397     check the file descriptor given to it, and always use coloring type
0398     according to C{ctype}.
0399 
0400     @param ctype: type of coloring
0401     @type ctype: string
0402     @param outdep: whether coloring depends on output file descriptor
0403     @type outdep: bool
0404     """
0405 
0406     _cglobals.outdep = outdep
0407     _cglobals.ctype = ctype
0408 
0409 
0410 class _Data: pass
0411 _cglobals = _Data()
0412 set_coloring_globals()
0413 
0414 # ========================================================================
0415 
0416 _color_packs = {}
0417 
0418 # ----------------------------------------
0419 # No coloring, all markup elements are just removed.
0420 
0421 _color_packs["none"] = (lambda c, s, p: s, lambda s: s, lambda r: r)
0422 
0423 
0424 # ----------------------------------------
0425 # ANSI terminal coloring.
0426 
0427 _term_head = "\033["
0428 _term_reset = "0;0m"
0429 _term_colors = {
0430     "bold": "01m",
0431     "underline": "04m",
0432     "black": "30m",
0433     "red": "31m",
0434     "green": "32m",
0435     "orange": "33m",
0436     "blue": "34m",
0437     "purple": "35m",
0438     "cyan": "36m",
0439     "grey": "37m",
0440 }
0441 
0442 def _color_term (col, seg, pcol):
0443 
0444     eseq = _term_colors.get(col)
0445     if eseq is not None:
0446         # If this segment is within another colored section,
0447         # repeat the outer color sequence at end, otherwise reset.
0448         # If outer and current colors match, do nothing.
0449         eseq2 = _term_reset
0450         peseq = _term_colors.get(pcol)
0451         if peseq:
0452             eseq2 = _term_head + peseq
0453         if eseq != eseq2:
0454             seg = _term_head + eseq + seg + _term_head + eseq2
0455     return seg
0456 
0457 
0458 _color_packs["term"] = (_color_term, lambda s: s, lambda r: r)
0459 
0460 
0461 # ----------------------------------------
0462 # HTML coloring.
0463 
0464 _html_colors = {
0465     "black": "#000000",
0466     "red": "#ff0000",
0467     "green": "#228b22",
0468     "orange": "#ff8040",
0469     "blue": "#0000ff",
0470     "purple": "#ff0080",
0471     "cyan": "#52f3ff",
0472     "grey": "#808080",
0473 }
0474 
0475 def _color_term (col, seg, pcol):
0476 
0477     if col == "bold":
0478         seg = "<b>%s</b>" % seg
0479     elif col == "underline":
0480         seg = "<u>%s</u>" % seg
0481     else:
0482         rgb = _html_colors.get(col)
0483         if rgb is not None:
0484             seg = "<font color='%s'>%s</font>" % (rgb, seg)
0485     return seg
0486 
0487 
0488 def _finalize_html (line):
0489 
0490     return line.replace("\n", "<br/>\n") + "<br/>"
0491 
0492 
0493 _color_packs["html"] = (_color_term, _escape_xml_ents, _finalize_html)
0494