Cleaner API, more comments, generator, toward py3k

Use simpler names for modules variables.
Detailled commets for each functions, with doctests.
Add a generator proxy.
Add a dirty argument parser that do not depend of argparse.
Separate args parser in private functions.
Clean the code.

colout.py

@@ -7,36 +7,69 @@
 
 import re
 
-styles = {"normal":0, "bold":1, "faint":2, "italic":3, "underline":4, "blink":5, "rapid_blink":6,
-"reverse":7, "conceal":8 }
-colors_mode8 = {"black":0, "red":1, "green":2, "yellow":3, "blue":4, "magenta":5, "cyan":6, "white":7}
-modes = {8:";", 256:";38;5;"}
 
+###########
+# Library #
+###########
 
-def colorin( text, color, style ):
+# Available styles
-    """Return the given text, surrounded by the given color ASCII markers."""
+styles = {
+    "normal":0, "bold":1, "faint":2, "italic":3, "underline":4,
+    "blink":5, "rapid_blink":6,
+    "reverse":7, "conceal":8
+}
+
+# Available color names in 8-colors mode
+colors = {
+    "black":0, "red":1, "green":2, "yellow":3, "blue":4,
+    "magenta":5, "cyan":6, "white":7
+}
+
+# Escaped end markers for given color modes
+endmarks = {8:";", 256:";38;5;"}
+
+def colorin( text, color = "red", style = "normal" ):
+    """
+    Return the given text, surrounded by the given color ASCII markers.
+
+    If the given color is a name that exists in available colors,
+    a 8-colors mode is assumed, else, a 256-colors mode.
+
+    The given style must exists in the available styles.
+
+    >>> colorin("Fetchez la vache", "red", "bold")
+    '\x1b[1;31mFetchez la vache\x1b[0m'
+    >>> colout.colorin("Faites chier la vache", 41, "normal")
+    '\x1b[0;38;5;41mFaites chier la vache\x1b[0m'
+    """
     # Special characters.
     start = "\033["
     stop = "\033[0m"
    
-    # Convert the color code.
+    # Convert the style code
-    cs = str(styles[style])
+    assert( style in styles )
+    style_code = str(styles[style])
 
     # 8 colors modes
-    if color in colors_mode8: 
+    if color in colors: 
         mode = 8
-        cc = str( 30 + colors_mode8[color] )
+        color_code = str( 30 + colors[color] )
     
     # 256 colors mode
     else: 
         mode = 256
-        cc = str( color ) 
+        color_nb = int( color )
+        assert( 0 <= color_nb <= 255 )
+        color_code = str( color ) 
 
-    return start + cs + modes[mode] + cc + "m" + text + stop
+    return start + style_code + endmarks[mode] + color_code + "m" + text + stop
 
 
-def colorout( text, match, prev_end, color, style, group=0 ):
+def colorout( text, match, prev_end, color = "red", style = "normal", group=0 ):
-    """Build the text from the previous match to the current one, coloring up the matching characters."""
+    """
+    Build the text from the previous re.match to the current one,
+    coloring up the matching characters.
+    """
     start = match.start(group)
     colored_text = text[prev_end:start]
     end = match.end(group)
@@ -45,17 +78,37 @@ def colorout( text, match, prev_end, color, style, group=0 ):
     return colored_text,end
 
 
-def colorup( text, pattern, color, style = "normal" ):
+def colorup( text, pattern, color = "red", style = "normal" ):
-    """Color up every characters that match the given patterns.
+    """
-    If groups are specified, only color up them and not the whole pattern."""
+    Color up every characters that match the given regexp patterns.
-    regex = re.compile(pattern, re.IGNORECASE)
+    If groups are specified, only color up them and not the whole pattern.
+
+    Colors and styles may be specified as a list of comma-separated values,
+    in which case the different matching groups may be formatted differently.
+    If there is less colors/styles than groups, the last format is used
+    for the additional groups.
+
+    >>> colorup("Fetchez la vache", "vache", "red", "bold")
+    'Fetchez la \x1b[1;31mvache\x1b[0m'
+    >>> colorup("Faites chier la vache", "[Fv]a", "red", "bold")
+    '\x1b[1;31mFa\x1b[0mites chier la \x1b[1;31mva\x1b[0mche'
+    >>> colorup("Faites Chier la Vache", "[A-Z](\S+)\s", "red", "bold")
+    'F\x1b[1;31maites\x1b[0m C\x1b[1;31mhier\x1b[0m la Vache'
+    >>> colorup("Faites Chier la Vache", "([A-Z])(\S+)\s", "red,green", "bold")
+    '\x1b[1;31mF\x1b[0m\x1b[1;32maites\x1b[0m \x1b[1;31mC\x1b[0m\x1b[1;32mhier\x1b[0m la Vache'
+    >>> colorup("Faites Chier la Vache", "([A-Z])(\S+)\s", "green")
+    '\x1b[0;32mF\x1b[0m\x1b[0;32maites\x1b[0m \x1b[0;32mC\x1b[0m\x1b[0;32mhier\x1b[0m la Vache'
+    >>> colorup("Faites Chier la Vache", "([A-Z])(\S+)\s", "blue", "bold,italic")
+    '\x1b[1;34mF\x1b[0m\x1b[3;34maites\x1b[0m \x1b[1;34mC\x1b[0m\x1b[3;34mhier\x1b[0m la Vache'
+    """
+    regex = re.compile(pattern)#, re.IGNORECASE)
 
     # Prepare the colored text.
     colored_text = ""
     end = 0
     for match in regex.finditer(text):
         
-        # If not groups are specified
+        # If no groups are specified
         if not match.groups():
             # Color the previous partial line,
             partial,end = colorout( text, match, end, color, style )
@@ -88,40 +141,121 @@ def colorup( text, pattern, color, style = "normal" ):
     return colored_text
 
 
-if __name__ == "__main__":
+def colorgen( items, pattern, color = "red", style = "normal" ):
-    import sys
+    """
-    import argparse
+    A generator that colors the items given in an iterable input.
 
+    >>> import math
+    >>> list(colorgen([str(i) for i in [math.pi,math.e]],"1","red"))
+    ['3.\x1b[0;31m1\x1b[0m4\x1b[0;31m1\x1b[0m59265359',
+     '2.7\x1b[0;31m1\x1b[0m828\x1b[0;31m1\x1b[0m82846']
+    """
+    for item in items:
+        yield colorup( item, pattern, color, style )
+
+
+######################
+# Command line tools #
+######################
+
+def __args_dirty__(argv,usage=""):
+    """
+    Roughly extract options from the command line arguments.
+    To be used only when argparse is not available.
+
+    Returns a tuple of (pattern,color,style,on_stderr).
+
+    >>> colout.__args_dirty__(["colout","pattern"],"usage")
+    ('pattern', 'red', 'normal', False)
+    >>> colout.__args_dirty__(["colout","pattern","colors","styles"],"usage")
+    ('pattern', 'colors', 'styles', False)
+    >>> colout.__args_dirty__(["colout","pattern","colors","styles","True"],"usage")
+    ('pattern', 'colors', 'styles', True)
+    """
+    import sys
+
+    # Use a dirty argument picker
+    # Check for bad usage or an help flag 
+    if    len(argv) < 2 \
+       or len(argv) > 5 \
+       or argv[1] == "--help" \
+       or argv[1] == "-h":
+        print(usage+"\n")
+        print("Usage:",argv[0],"<pattern> <color(s)> [<style(s)>] [<print on stderr?>]")
+        print("\tAvailable colors:"," ".join(colors))
+        print("\tAvailable styles:"," ".join(styles))
+        print("Example:",argv[0],"'^(def)\s+(\w*).*$' blue,magenta italic,bold < colout.py")
+        sys.exit(1)
+    
+    assert( len(argv) >= 2 )
+    # Get mandatory arguments
+    pattern = argv[1]
+
+    # default values for optional args
+    color = "red"
+    style = "normal"
+    on_stderr = False
+
+    if len(argv) >= 3:
+        color = argv[2]
+        if len(argv) >= 4:
+            style = argv[3]
+            if len(argv) == 5:
+                on_stderr = bool(argv[4])
+    
+    return pattern,color,style,on_stderr
+
+
+def __args_parse__(argv,usage=""):
+    """
+    Parse command line arguments with the argparse library.
+    Returns a tuple of (pattern,color,style,on_stderr).
+    """
     parser = argparse.ArgumentParser(
-        description="A regular expression based formatter that color up an arbitrary text output stream.")
+        description=usage)
 
     parser.add_argument("pattern", metavar="REGEX", type=str, nargs=1,
             help="A regular expression")
 
     parser.add_argument("color", metavar="COLOR", type=str, nargs='?',
             default="red",
-            help="A number in [0…255] or one of the following colors: "+" ".join(colors_mode8) )
+            help="A number in [0…255], one of the available colors or a comma-separated list of values. \
+                        Available colors: "+" ".join(colors) )
 
     parser.add_argument("style", metavar="STYLE", type=str, nargs='?',
             default="bold",
-            help="One of the following styles: "+" ".join(styles) )
+            help="One of the available styles or a comma-separated list of styles.\
+                        Available styles: "+" ".join(styles) )
 
     parser.add_argument("-e", "--stderr", action="store_true",
             help="Output on the stderr instead of stdout")
 
     args = parser.parse_args()
 
-    while True:
+    return args.pattern[0], args.color, args.style, args.stderr
-        line = sys.stdin.readline()
-        if line == '':
-            break
-        try:
-            if not args.stderr:
-                print colorup( line, args.pattern[0], args.color, args.style ),
-                sys.stdout.flush()
-            else:
-                print >> sys.stderr, colorup( line, args.pattern[0], args.color, args.style ),
-                sys.stderr.flush()
-        except:
-            pass
 
+
+if __name__ == "__main__":
+    import sys
+
+    usage="A regular expression based formatter that color up an arbitrary text stream."
+
+    try:
+        import argparse
+
+    # if argparse is not installed
+    except ImportError:
+        pattern,color,style,on_stderr = __args_dirty__(sys.argv,usage)
+        
+    # if argparse is available
+    else:
+        pattern,color,style,on_stderr = __args_parse__(sys.argv,usage)
+
+    # use the generator: output lines as they come
+    for colored in colorgen( sys.stdin, pattern, color, style ):
+        if on_stderr:
+            sys.stderr.write(colored)
+            sys.stderr.flush()
+        else:
+            sys.stdout.write(colored)
+            sys.stdout.flush()