from __future__ import absolute_import, division

import re

import numba
from numba import unittest_support as unittest
from numba.compiler import compile_isolated, Flags
from numba import types
from numba.io_support import StringIO

try:
    import jinja2
except ImportError:
    jinja2 = None


@unittest.skipIf(jinja2 is None, "please install the 'jinja2' package")
class TestAnnotation(unittest.TestCase):

    def test_exercise_code_path(self):
        """
        Ensures template.html is available
        """

        def foo(n, a):
            s = a
            for i in range(n):
                s += i
            return s

        cres = compile_isolated(foo, [types.int32, types.int32])
        ta = cres.type_annotation
        self.assertIs(ta.html_output, None)

        buf = StringIO()
        ta.html_annotate(buf)
        output = buf.getvalue()
        buf.close()
        self.assertIn("foo", output)

    def test_exercise_code_path_with_lifted_loop(self):
        """
        Ensures that lifted loops are handled correctly in obj mode
        """
        # the functions to jit
        def bar(x):
            return x

        def foo(x):
            h = 0.
            for k in range(x):
                h = h + k
            if x:
                h = h - bar(x)
            return h

        # compile into an isolated context
        flags = Flags()
        flags.set('enable_pyobject')
        flags.set('enable_looplift')
        cres = compile_isolated(foo, [types.intp], flags=flags)

        ta = cres.type_annotation
        self.assertIs(ta.html_output, None)

        buf = StringIO()
        ta.html_annotate(buf)
        output = buf.getvalue()
        buf.close()
        self.assertIn("bar", output)
        self.assertIn("foo", output)
        self.assertIn("LiftedLoop", output)

if __name__ == '__main__':
    unittest.main()

@@ -1817,16 +1817,1928 @@
 utput)%0A%0A
+    def test_html_output_with_lifted_loop(self):%0A        %22%22%22%0A        Test some format and behavior of the html annotation with lifted loop%0A        %22%22%22%0A        @numba.jit%0A        def udt(x):%0A            object()  # to force object mode%0A            z = 0%0A            for i in range(x):  # this line is tagged%0A                z += i%0A            return z%0A%0A        # Regex pattern to check for the %22lifted_tag%22 in the line of the loop%0A        re_lifted_tag = re.compile(%0A            r'%3Ctd class=%22lifted_tag%22%3E%5Cs*'%0A            r'%5B&nbsp;%5D+for i in range%5C(x%5C):  # this line is tagged%5Cs*'%0A            r'%3C/td%3E', re.MULTILINE)%0A%0A        # Compile int64 version%0A        sig_i64 = (types.int64,)%0A        udt.compile(sig_i64)  # compile with lifted loop%0A        cres = udt.overloads%5Bsig_i64%5D%0A%0A        # Make html output%0A        buf = StringIO()%0A        cres.type_annotation.html_annotate(buf)%0A        output = buf.getvalue()%0A        buf.close()%0A%0A        # There should be only one function output.%0A        self.assertEqual(output.count(%22Function name: udt%22), 1)%0A        self.assertEqual(output.count(%22with signature: %7B%7D%22.format(sig_i64)), 1)%0A        # Ensure the loop is tagged%0A        self.assertEqual(len(re.findall(re_lifted_tag, output)), 1)%0A%0A        # Compile float64 version%0A        sig_f64 = (types.float64,)%0A        udt.compile(sig_f64)%0A        cres = udt.overloads%5Bsig_f64%5D%0A%0A        # Make html output%0A        buf = StringIO()%0A        cres.type_annotation.html_annotate(buf)%0A        output = buf.getvalue()%0A        buf.close()%0A%0A        # There should be two function output%0A        self.assertEqual(output.count(%22Function name: udt%22), 2)%0A        self.assertEqual(output.count(%22with signature: %7B%7D%22.format(sig_i64)), 1)%0A        self.assertEqual(output.count(%22with signature: %7B%7D%22.format(sig_f64)), 1)%0A        # Ensure the loop is tagged in both output%0A        self.assertEqual(len(re.findall(re_lifted_tag, output)), 2)%0A%0A%0A
 if __nam

Test the output of html annotate with lifted loop