Refactor parser code, add filename option

jinja2_highlight/highlight.py

@@ -5,7 +5,10 @@
 from jinja2.ext import Extension, Markup
 
 from pygments import highlight
-from pygments.lexers import get_lexer_by_name, guess_lexer
+from pygments.lexers import(
+    get_lexer_by_name, guess_lexer, guess_lexer_for_filename
+)
+from pygments.lexers.special import TextLexer
 from pygments.formatters import HtmlFormatter
 from pygments.util import ClassNotFound
 
@@ -23,6 +26,18 @@ class HighlightExtension(Extension):
 
         {% endhighlight %}
 
+    If you are unable to provide the language to the highlight extension but can
+    provide the filename, use the filename option.
+
+        {% highlight filename="beer.py" %}
+
+        from fridge import Beer
+
+        pint_glass = Beer()
+        pint_glass.drink()
+
+        {% endhighlight %}
+
     Line numbers can be turned on if it is assigned. Either 'inline' or 'table'
     settings in Pygments can be used.
 
@@ -50,58 +65,38 @@ class HighlightExtension(Extension):
     def parse(self, parser):
         lineno = next(parser.stream).lineno
 
-        # extract the language and line numbering setting if available
-        if not parser.stream.current.test('block_end'):
-            # If first up we have an assignment, e.g. lineno='inline', work with that
-            if parser.stream.current.type == 'name':
+        # NOTE: The _highlight function parameter order must match the order of arg_list
+        arg_list = ['lang', 'lineno', 'filename']
+        parsed_args = {}
+
+        while not parser.stream.current.test('block_end'):
+            this_token_type = parser.stream.current.type
+
+            # If up we have an assignment, e.g. lineno='inline', work with that
+            if this_token_type == 'name':
                 name = parser.stream.expect('name')
-                # If the assign is lineno
-                if name.value == 'lineno':
-                    if parser.stream.skip_if('assign'):
-                        # Assume no language and then add the assigned line number setting
-                        args = [nodes.Const(None)]
-                        args.append(parser.parse_expression())
-                # If it's not a lineno assignment, ignore it
+                if name.value in arg_list and parser.stream.skip_if('assign'):
+                    parsed_args[name.value] = parser.parse_expression()
                 else:
-                    if parser.stream.skip_if('assign'):
-                        next(parser.stream)
-                        # Set our language and line number setting as None
-                        args = [nodes.Const(None), nodes.Const(None)]
+                    parser.fail('Unrecognized argument: %s' % name.value)
+            elif this_token_type == 'string':
+                # The only valid string literal argument is the language itself
+                parsed_args['lang'] = parser.parse_expression()
             else:
-                # Otherwise if our first item is not a line numbering setting,
-                # assume it's the language setting
-                args = [parser.parse_expression()]
-
-                # If we have a comma next
-                if parser.stream.skip_if('comma'):
-                    # Check to see if we have a lineno assignment
-                    if parser.stream.current.type == 'name':
-                        name = parser.stream.expect('name')
-                        if name.value == 'lineno':
-                            if parser.stream.skip_if('assign'):
-                                args.append(parser.parse_expression())
-                        # If the name of the variable being assigned is not lineno
-                        # ignore it
-                        else:
-                            if parser.stream.skip_if('assign'):
-                                next(parser.stream)
-                                args.append(nodes.Const(None))
-                # Otherwise if there's nothing after the language, set the
-                # line number setting as None
-                else:
-                    args.append(nodes.Const(None))
-        else:
-            # Otherwise if there are no additional arguments, set lang
-            # and line numbering to None
-            args = [nodes.Const(None), nodes.Const(None)]
+                parser.fail('Unexpected %s encountered' % this_token_type)
+
+            # Skip over optional commas
+            parser.stream.skip_if('comma')
+
+        args = [parsed_args.get(a, nodes.Const(None)) for a in arg_list]
 
-        # body of the block
+        # body of the block (the source code we want to highlight)
         body = parser.parse_statements(['name:endhighlight'], drop_needle=True)
 
         return nodes.CallBlock(self.call_method('_highlight', args),
                                [], [], body).set_lineno(lineno)
 
-    def _highlight(self, lang, linenos, caller=None):
+    def _highlight(self, lang, linenos, filename, caller=None):
         # highlight code using Pygments
         body = caller()
 
@@ -113,13 +108,15 @@ def _highlight(self, lang, linenos, caller=None):
             cssclass = None
 
         try:
-            if lang is None:
-                lexer = guess_lexer(body)
-            else:
+            if lang is not None:
                 lexer = get_lexer_by_name(lang, stripall=False)
+            elif filename is not None:
+                lexer = guess_lexer_for_filename(filename, body)
+            else:
+                lexer = guess_lexer(body)
         except ClassNotFound as e:
-            print(e)
-            sys.exit(1)
+            # default to the plaintext lexer
+            lexer = TextLexer()
 
         # Set the cssclass if we have one
         # The linenos setting expects either 'inline' or 'table', as per Pygment's

tests/test_highlight.py

@@ -6,6 +6,12 @@
 
 class HighlightExtensionTestCase(unittest.TestCase):
 
+    def assertHtmlListEqual(self, a, b):
+        # Normalize the HTML lists so white space doesn't cause a failure.
+        html_a = "".join(a)
+        html_b = "".join(b)
+        self.assertEqual(html_a, html_b)
+
     rendered = [
                 u'<div',
                 u'class="highlight"><pre>',
@@ -26,7 +32,17 @@ def test_python_tpl(self):
             {% endhighlight %}
         ''')
 
-        assert tpl.render().split() == self.rendered
+        self.assertHtmlListEqual(tpl.render().split(), self.rendered)
+
+    def test_python_tpl_by_filename(self):
+        env = Environment(extensions=['jinja2_highlight.HighlightExtension'])
+        tpl = env.from_string('''
+            {% highlight filename="hello.py" %}
+               print("Hello world")
+            {% endhighlight %}
+        ''')
+
+        self.assertHtmlListEqual(tpl.render().split(), self.rendered)
 
     def test_python_tpl_with_autoescape(self):
         # See: https://github.com/tlatsas/jinja2-highlight/pull/1
@@ -40,7 +56,7 @@ def test_python_tpl_with_autoescape(self):
             {% endautoescape %}
         ''')
 
-        assert tpl.render().split() == self.rendered
+        self.assertHtmlListEqual(tpl.render().split(), self.rendered)
 
     inline_rendered = [
                 u'<div',
@@ -64,7 +80,7 @@ def test_python_tpl_with_inline(self):
             {% endhighlight %}
         ''')
 
-        assert tpl.render().split() == self.inline_rendered
+        self.assertHtmlListEqual(tpl.render().split(), self.inline_rendered)
 
     table_rendered = [
                 u'<table',
@@ -91,57 +107,62 @@ def test_python_tpl_with_table(self):
             {% endhighlight %}
         ''')
 
-        assert tpl.render().split() == self.table_rendered
+        self.assertHtmlListEqual(tpl.render().split(), self.table_rendered)
 
     inline_no_lang_rendered = [
-                u'<div',
-                u'class="highlight"><pre><span',
-                u'class="lineno">1</span>',
-                u'<span',
-                u'class="n">print</span><span',
-                u'class="p">(</span><span',
-                u'class="s">&quot;Hello',
-                u'world&quot;</span><span',
-                u'class="p">)</span>',
-                u'</pre></div>'
-            ]
+        u'<div',
+        u'class="highlight"><pre><span',
+        u'class="lineno">1',
+        u'</span>',
+        u'<span',
+        u'class="kn">from</span>',
+        u'<span',
+        u'class="nn">mypackage.mymodule</span>',
+        u'<span',
+        u'class="kn">import</span>',
+        u'<span',
+        u'class="n">myfn</span>',
+        u'</pre></div>'
+    ]
 
     def test_python_tpl_with_inline_no_lang(self):
         env = Environment(extensions=['jinja2_highlight.HighlightExtension'])
         tpl = env.from_string('''
             {% highlight lineno="inline" %}
-               print("Hello world")
+               from mypackage.mymodule import myfn
             {% endhighlight %}
         ''')
 
-        assert tpl.render().split() == self.inline_no_lang_rendered
+        self.assertHtmlListEqual(tpl.render().split(), self.inline_no_lang_rendered)
 
     table_no_lang_rendered = [
-                u'<table',
-                u'class="highlighttable"><tr><td',
-                u'class="linenos"><div',
-                u'class="linenodiv"><pre>1</pre></div></td><td',
-                u'class="code"><div',
-                u'class="highlight"><pre>',
-                u'<span',
-                u'class="n">print</span><span',
-                u'class="p">(</span><span',
-                u'class="s">&quot;Hello',
-                u'world&quot;</span><span',
-                u'class="p">)</span>',
-                u'</pre></div>',
-                u'</td></tr></table>'
-            ]
+        u'<table',
+        u'class="highlighttable"><tr><td',
+        u'class="linenos"><div',
+        u'class="linenodiv"><pre>1</pre></div></td><td',
+        u'class="code"><div',
+        u'class="highlight"><pre>',
+        u'<span',
+        u'class="kn">from</span>',
+        u'<span',
+        u'class="nn">mypackage.mymodule</span>',
+        u'<span',
+        u'class="kn">import</span>',
+        u'<span',
+        u'class="n">myfn</span>',
+        u'</pre></div>',
+        u'</td></tr></table>'
+    ]
 
     def test_python_tpl_with_table_no_lang(self):
         env = Environment(extensions=['jinja2_highlight.HighlightExtension'])
         tpl = env.from_string('''
             {% highlight lineno="table" %}
-               print("Hello world")
+               from mypackage.mymodule import myfn
             {% endhighlight %}
         ''')
 
-        assert tpl.render().split() == self.table_no_lang_rendered
+        self.assertHtmlListEqual(tpl.render().split(), self.table_no_lang_rendered)
 
     cssclass_rendered = [
                 u'<div',
@@ -165,4 +186,23 @@ def test_python_tpl_with_cssclass(self):
             {% endhighlight %}
         ''')
 
-        assert tpl.render().split() == self.cssclass_rendered
+        self.assertHtmlListEqual(tpl.render().split(), self.cssclass_rendered)
+
+    plaintext_rendered = [
+        u'<div',
+        u'class="highlight"><pre>',
+        u'abcdefg',
+        u'</pre></div>'
+    ]
+
+    def test_unrecognized_language_defaults_to_plaintext(self):
+        env = Environment(extensions=['jinja2_highlight.HighlightExtension'])
+
+        # Unrecognized language
+        tpl = env.from_string('''
+            {% highlight %}
+               abcdefg
+            {% endhighlight %}
+        ''')
+
+        self.assertHtmlListEqual(tpl.render().split(), self.plaintext_rendered)