#!/usr/bin/python

"""

CodeHilite Extension for Python-Markdown

========================================

Adds code/syntax highlighting to standard Python-Markdown code blocks.

Copyright 2006-2008 [Waylan Limberg](http://achinghead.com/).

Project website: <http://www.freewisdom.org/project/python-markdown/CodeHilite>

Contact: markdown@freewisdom.org

License: BSD (see ../docs/LICENSE for details)

Dependencies:

* [Python 2.3+](http://python.org/)

* [Markdown 2.0+](http://www.freewisdom.org/projects/python-markdown/)

* [Pygments](http://pygments.org/)

"""

import markdown

# ------------------ The Main CodeHilite Class ----------------------

class CodeHilite:

    """

    Determine language of source code, and pass it into the pygments hilighter.

    Basic Usage:

        >>> code = CodeHilite(src = 'some text')

        >>> html = code.hilite()

    * src: Source string or any object with a .readline attribute.

    * linenos: (Boolen) Turn line numbering 'on' or 'off' (off by default).

    * css_class: Set class name of wrapper div ('codehilite' by default).

    Low Level Usage:

        >>> code = CodeHilite()

        >>> code.src = 'some text' # String or anything with a .readline attr.

        >>> code.linenos = True  # True or False; Turns line numbering on or of.

        >>> html = code.hilite()

    """

    def __init__(self, src=None, linenos=False, css_class="codehilite",

                lang=None, style='default', noclasses=False, tab_length=4):

        self.src = src

        self.lang = lang

        self.linenos = linenos

        self.css_class = css_class

        self.style = style

        self.noclasses = noclasses

        self.tab_length = tab_length

    def hilite(self):

        """

        Pass code to the [Pygments](http://pygments.pocoo.org/) highliter with

        optional line numbers. The output should then be styled with css to

        your liking. No styles are applied by default - only styling hooks

        (i.e.: <span class="k">).

        returns : A string of html.

        """

        self.src = self.src.strip('\n')

        if self.lang == None:

            self._getLang()

        try:

            from pygments import highlight

            from pygments.lexers import get_lexer_by_name, guess_lexer, \

                                        TextLexer

            from pygments.formatters import HtmlFormatter

        except ImportError:

            # just escape and pass through

            txt = self._escape(self.src)

            if self.linenos:

                txt = self._number(txt)

            else :

                txt = '<div class="%s"><pre>%s</pre></div>\n'% \

                        (self.css_class, txt)

            return txt

        else:

            try:

                lexer = get_lexer_by_name(self.lang)

            except ValueError:

                try:

                    lexer = guess_lexer(self.src)

                except ValueError:

                    lexer = TextLexer()

            formatter = HtmlFormatter(linenos=self.linenos,

                                      cssclass=self.css_class,

                                      style=self.style,

                                      noclasses=self.noclasses)

            return highlight(self.src, lexer, formatter)

    def _escape(self, txt):

        """ basic html escaping """

        txt = txt.replace('&', '&')

        txt = txt.replace('<', '&lt;')

        txt = txt.replace('>', '>')

        txt = txt.replace('"', '"')

        return txt

    def _number(self, txt):

        """ Use <ol> for line numbering """

        # Fix Whitespace

        txt = txt.replace('\t', ' '*self.tab_length)

        txt = txt.replace(" "*4, "    ")

        txt = txt.replace(" "*3, "   ")

        txt = txt.replace(" "*2, "  ")

        # Add line numbers

        lines = txt.splitlines()

        txt = '<div class="codehilite"><pre><ol>\n'

        for line in lines:

            txt += '\t<li>%s</li>\n'% line

        txt += '</ol></pre></div>\n'

        return txt

    def _getLang(self):

        """

        Determines language of a code block from shebang lines and whether said

        line should be removed or left in place. If the sheband line contains a

        path (even a single /) then it is assumed to be a real shebang lines and

        left alone. However, if no path is given (e.i.: #!python or :::python)

        then it is assumed to be a mock shebang for language identifitation of a

        code fragment and removed from the code block prior to processing for

        code highlighting. When a mock shebang (e.i: #!python) is found, line

        numbering is turned on. When colons are found in place of a shebang

        (e.i.: :::python), line numbering is left in the current state - off

        by default.

        """

        import re

        #split text into lines

        lines = self.src.split("\n")

        #pull first line to examine

        fl = lines.pop(0)

        c = re.compile(r'''

            (?:(?:::+)|(?P<shebang>[#]!))	# Shebang or 2 or more colons.

            (?P<path>(?:/\w+)*[/ ])?        # Zero or 1 path

            (?P<lang>[\w+-]*)               # The language

            ''',  re.VERBOSE)

        # search first line for shebang

        m = c.search(fl)

        if m:

            # we have a match

            try:

                self.lang = m.group('lang').lower()

            except IndexError:

                self.lang = None

            if m.group('path'):

                # path exists - restore first line

                lines.insert(0, fl)

            if m.group('shebang'):

                # shebang exists - use line numbers

                self.linenos = True

        else:

            # No match

            lines.insert(0, fl)

        self.src = "\n".join(lines).strip("\n")

# ------------------ The Markdown Extension -------------------------------

class HiliteTreeprocessor(markdown.treeprocessors.Treeprocessor):

    """ Hilight source code in code blocks. """

    def run(self, root):

        """ Find code blocks and store in htmlStash. """

        blocks = root.getiterator('pre')

        for block in blocks:

            children = block.getchildren()

            if len(children) == 1 and children[0].tag == 'code':

                code = CodeHilite(children[0].text,

                            linenos=self.config['force_linenos'][0],

                            css_class=self.config['css_class'][0],

                            style=self.config['pygments_style'][0],

                            noclasses=self.config['noclasses'][0],

                            tab_length=self.markdown.tab_length)

                placeholder = self.markdown.htmlStash.store(code.hilite(),

                                                            safe=True)

                # Clear codeblock in etree instance

                block.clear()

                # Change to p element which will later

                # be removed when inserting raw html

                block.tag = 'p'

                block.text = placeholder

class CodeHiliteExtension(markdown.Extension):

    """ Add source code hilighting to markdown codeblocks. """

    def __init__(self, configs):

        # define default configs

        self.config = {

            'force_linenos' : [False, "Force line numbers - Default: False"],

            'css_class' : ["codehilite",

                           "Set class name for wrapper <div> - Default: codehilite"],

            'pygments_style' : ['default', 'Pygments HTML Formatter Style (Colorscheme) - Default: default'],

            'noclasses': [False, 'Use inline styles instead of CSS classes - Default false']

            }

        # Override defaults with user settings

        for key, value in configs:

            self.setConfig(key, value)

    def extendMarkdown(self, md, md_globals):

        """ Add HilitePostprocessor to Markdown instance. """

        hiliter = HiliteTreeprocessor(md)

        hiliter.config = self.config

        md.treeprocessors.add("hilite", hiliter, "_begin")

        md.registerExtension(self)

def makeExtension(configs={}):

  return CodeHiliteExtension(configs=configs)