pygments ftw

pygmenter.py

#!/usr/bin/python
#

import os
import sys
import argparse

import pygments
import pygments.lexers
import pygments.formatters

def readfile(filename, start = None, end = None):
    i = 0
    buf = ''
    for ln in file(filename):
        i += 1
        if start is not None and i < start:
            continue
        if end is not None and i >= end:
            break
        print i, ln,
        buf += ln
    return buf

def snip_values(args):
    snip = args.snip
    snip = snip.split(':')
    start = end = None
    try:
        start = int(snip[0])
    except:
        pass
    try:
        end = int(snip[1])
    except:
        pass
    return {'start' : start, 'end' : end } 

def pygmentize(intext, args, linenostart):
    if linenostart is None: linenostart = 1
    lexer = pygments.lexers.get_lexer_by_name(args.lexer)
    fmtr = pygments.formatters.HtmlFormatter(
        linenos='inline', 
        linenostart=linenostart, 
        stripall=True)
    return pygments.highlight(intext, lexer, fmtr)

def _snip_to_text(snip):
    snip = snip.copy()
    if len([x for x in snip.values() if x is not None]) == 0:
        return ''
    if snip['start'] is None:
        snip['start'] = 'START'
    if snip['end'] is None:
        snip['end'] = 'END'
    return '_%(start)s-%(end)s' % snip

def outputfile(infilename, outtext, snip):
    outfile = os.path.split(infilename)[-1]
    sniptext = _snip_to_text(snip)
    outfile = '%s%s.html' % (outfile, sniptext)
    fd = file(outfile, 'w')
    fd.write(outtext)
    fd.flush()
    fd.close()

def parse_arguments(description):
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument(
        'filename',
        nargs=1,
        help='name of file to highlight.')
    parser.add_argument(
        '--snip',
        type=str,
        default=":",
        help='Only read these line numbers. Form is "start:end", with'
            ' either side possibly being blank, which implies beginning/end'
            ' of the file, respectively. So "5:10" means lines 5-10 inclusive,'
            ' and ":10" means beginning of file to line 10 inclusive.' 
            ' default: "%(default)s"')
    parser.add_argument(
        '-l', '--lexer', 
        type=str,
        default='python', 
        help='pygments lexer name. See: http://pygments.org/docs/lexers/.'
            ' default: "%(default)s"')

    return parser.parse_args()

def main(args):
    snip = snip_values(args)
    filename = os.path.abspath(args.filename[0])
    intext = readfile(filename, **snip)
    outtext = pygmentize(intext, args, linenostart=snip['start'])
    outputfile(filename, outtext, snip)
    return 0

USAGE = '''
Default pygmentize call which does
-O linenos=inline and creates 
and {filename}.html file. 

'''

if __name__ == '__main__':
    args = parse_arguments(USAGE)
    v = main(args)
    sys.exit(v)