/kai/lib/highlight.py

 1from docutils.parsers import rst
 2from pygments import highlight
 3from pygments.lexers import get_lexer_by_name, get_all_lexers
 4from pygments.formatters import HtmlFormatter
 5
 6from kai.lib.pygmentsupport import code_block
 7
 8__all__ = ['code_highlight']
 9
10langdict = {}
11for lang in get_all_lexers():
12    langdict[lang[1][0]] = lang[0]
13
14formatter = HtmlFormatter(linenos=True, cssclass="syntax", encoding='utf-8')
15
16def code_highlight(code, truncate_lines=None):
17    # Ensure that code-block is regged properly
18    rst.directives.register_directive('code-block', code_block)
19    source = code.code
20    if truncate_lines:
21        split_source = source.split('\n')
22        if len(split_source) > truncate_lines:
23            source = split_source[:truncate_lines-1]
24            source.append('...')
25            source = ''.join(source)
26    lexer = get_lexer_by_name(code.language, stripall=True)
27    
28    return highlight(source, lexer, formatter).decode('utf-8')