1passでpygmentsの出力したhtmlにスタイルを埋め込む

前作ったやつは、スタイル出力、HTML出力、スタイル埋め込みという手順を踏む必要があって面倒なので、pygmentsのAPIを使って一発で処理するようにしてみた。

前のは出力されたcssとHTMLを処理するのでstyle, lexerなんかは好きにすればよかったけど、
今回は1passなのでそこら辺の指定をできるようにしてみた。



from pygments import highlight
from pygments.lexers import get_lexer_by_name, get_lexer_for_filename
from pygments.formatters import get_formatter_by_name
from xml.etree.ElementTree import ElementTree, XML
import cssutils
from cssutils.css import CSSRule
from sys import stdout
from optparse import OptionParser

def write(filename, style_name='colorful', formatter_name='html', lexer_name=None, out=stdout):
    fmtter = get_formatter_by_name(formatter_name, style=style_name)
    if lexer_name:
        lexer = get_lexer_by_name(lexer_name)
    else:
        lexer = get_lexer_for_filename(filename)
    css = cssutils.parseString(fmtter.get_style_defs())
    rules = dict([(x.selectorText, x.style.cssText.replace('\n', '')) for x in css.cssRules if x.type == CSSRule.STYLE_RULE])
    tree = ElementTree(XML(highlight(open(filename).read(), lexer, fmtter)))
    for elem in tree.getiterator():
        c = elem.get('class', None)
        if c:
            style = rules.get('.'+c, None)
            if style:
                elem.attrib['style'] = style
    tree.write(out)

if __name__ == '__main__':
    optparser = OptionParser(usage=u'usage: %prog [options] file')
    optparser.add_option('-S', '--style', dest='style', help='style', default='colorful')
    optparser.add_option('-f', '--formatter', dest='formatter', help='formatter', default='html')
    optparser.add_option('-L', '--lexer', dest='lexer', help='lexer')
    options, args = optparser.parse_args()
    if len(args) == 1:
        write(args[0], options.style, options.formatter, options.lexer, stdout)
    else:
        print '** no input file specified **'
        optparser.print_help()