refactor visit_Output

* `finalize` is generated once and cached for all nodes.
* Extract common behavior for native env.

Removed the compiler behavior where groups of nodes would generate a
format string. Instead, individual nodes are always yielded. This made
rendering 30% faster in the examples, and simplifies the code. It also
removes the issue where Python would report either the first or last
line of the multi-line format expression, messing up the traceback line
number mapping.

5 files changed, 188 insertions, 297 deletions

diff --git a/jinja2/compiler.py b/jinja2/compiler.py
index 50e00ab..3c69df8 100644
--- a/jinja2/compiler.py
+++ b/jinja2/compiler.py
@@ -8,8 +8,8 @@
     :copyright: (c) 2017 by the Jinja Team.
     :license: BSD, see LICENSE for more details.
 """
+from collections import namedtuple
 from itertools import chain
-from copy import deepcopy
 from keyword import iskeyword as is_python_keyword
 from functools import update_wrapper
 from jinja2 import nodes
@@ -1221,75 +1221,136 @@ class CodeGenerator(NodeVisitor):
         self.newline(node)
         self.visit(node.node, frame)
 
-    def visit_Output(self, node, frame):
-        # if we have a known extends statement, we don't output anything
-        # if we are in a require_output_check section
-        if self.has_known_extends and frame.require_output_check:
-            return
+    _FinalizeInfo = namedtuple("_FinalizeInfo", ("const", "src"))
+    #: The default finalize function if the environment isn't configured
+    #: with one. Or if the environment has one, this is called on that
+    #: function's output for constants.
+    _default_finalize = text_type
+    _finalize = None
+
+    def _make_finalize(self):
+        """Build the finalize function to be used on constants and at
+        runtime. Cached so it's only created once for all output nodes.
+
+        Returns a ``namedtuple`` with the following attributes:
+
+        ``const``
+            A function to finalize constant data at compile time.
+
+        ``src``
+            Source code to output around nodes to be evaluated at
+            runtime.
+        """
+        if self._finalize is not None:
+            return self._finalize
 
-        finalize = text_type
-        finalize_src = None
-        allow_constant_finalize = True
+        finalize = default = self._default_finalize
+        src = None
 
         if self.environment.finalize:
+            src = "environment.finalize("
             env_finalize = self.environment.finalize
-            finalize_src = "environment.finalize("
 
             def finalize(value):
-                return text_type(env_finalize(value))
+                return default(env_finalize(value))
 
             if getattr(env_finalize, "contextfunction", False):
-                finalize_src += "context, "
-                allow_constant_finalize = False
+                src += "context, "
+                finalize = None
             elif getattr(env_finalize, "evalcontextfunction", False):
-                finalize_src += "context.eval_ctx, "
-                allow_constant_finalize = False
+                src += "context.eval_ctx, "
+                finalize = None
             elif getattr(env_finalize, "environmentfunction", False):
-                finalize_src += "environment, "
+                src += "environment, "
 
                 def finalize(value):
-                    return text_type(env_finalize(self.environment, value))
+                    return default(env_finalize(self.environment, value))
+
+        self._finalize = self._FinalizeInfo(finalize, src)
+        return self._finalize
+
+    def _output_const_repr(self, group):
+        """Given a group of constant values converted from ``Output``
+        child nodes, produce a string to write to the template module
+        source.
+        """
+        return repr(concat(group))
+
+    def _output_child_to_const(self, node, frame, finalize):
+        """Try to optimize a child of an ``Output`` node by trying to
+        convert it to constant, finalized data at compile time.
+
+        If :exc:`Impossible` is raised, the node is not constant and
+        will be evaluated at runtime. Any other exception will also be
+        evaluated at runtime for easier debugging.
+        """
+        const = node.as_const(frame.eval_ctx)
+
+        if frame.eval_ctx.autoescape:
+            const = escape(const)
 
-        # if we are inside a frame that requires output checking, we do so
-        outdent_later = False
+        # Template data doesn't go through finalize.
+        if isinstance(node, nodes.TemplateData):
+            return text_type(const)
+
+        return finalize.const(const)
+
+    def _output_child_pre(self, node, frame, finalize):
+        """Output extra source code before visiting a child of an
+        ``Output`` node.
+        """
+        if frame.eval_ctx.volatile:
+            self.write("(escape if context.eval_ctx.autoescape else to_string)(")
+        elif frame.eval_ctx.autoescape:
+            self.write("escape(")
+        else:
+            self.write("to_string(")
+
+        if finalize.src is not None:
+            self.write(finalize.src)
+
+    def _output_child_post(self, node, frame, finalize):
+        """Output extra source code after visiting a child of an
+        ``Output`` node.
+        """
+        self.write(")")
+
+        if finalize.src is not None:
+            self.write(")")
+
+    def visit_Output(self, node, frame):
+        # If an extends is active, don't render outside a block.
         if frame.require_output_check:
-            self.writeline('if parent_template is None:')
+            # A top-level extends is known to exist at compile time.
+            if self.has_known_extends:
+                return
+
+            self.writeline("if parent_template is None:")
             self.indent()
-            outdent_later = True
 
-        # try to evaluate as many chunks as possible into a static
-        # string at compile time.
+        finalize = self._make_finalize()
         body = []
+
+        # Evaluate constants at compile time if possible. Each item in
+        # body will be either a list of static data or a node to be
+        # evaluated at runtime.
         for child in node.nodes:
             try:
-                # If the finalize function needs context, and this isn't
-                # template data, evaluate the node at render.
                 if not (
-                    allow_constant_finalize
+                    # If the finalize function requires runtime context,
+                    # constants can't be evaluated at compile time.
+                    finalize.const
+                    # Unless it's basic template data that won't be
+                    # finalized anyway.
                     or isinstance(child, nodes.TemplateData)
                 ):
                     raise nodes.Impossible()
 
-                const = child.as_const(frame.eval_ctx)
-            except nodes.Impossible:
-                body.append(child)
-                continue
-
-            # the frame can't be volatile here, becaus otherwise the
-            # as_const() function would raise an Impossible exception
-            # at that point.
-            try:
-                if frame.eval_ctx.autoescape:
-                    const = escape(const)
-
-                # Only call finalize on expressions, not template data.
-                if isinstance(child, nodes.TemplateData):
-                    const = text_type(const)
-                else:
-                    const = finalize(const)
-            except Exception:
-                # if something goes wrong here we evaluate the node
-                # at runtime for easier debugging
+                const = self._output_child_to_const(child, frame, finalize)
+            except (nodes.Impossible, Exception):
+                # The node was not constant and needs to be evaluated at
+                # runtime. Or another error was raised, which is easier
+                # to debug at runtime.
                 body.append(child)
                 continue
 
@@ -1298,79 +1359,42 @@ class CodeGenerator(NodeVisitor):
             else:
                 body.append([const])
 
-        # if we have less than 3 nodes or a buffer we yield or extend/append
-        if len(body) < 3 or frame.buffer is not None:
-            if frame.buffer is not None:
-                # for one item we append, for more we extend
-                if len(body) == 1:
-                    self.writeline('%s.append(' % frame.buffer)
+        if frame.buffer is not None:
+            if len(body) == 1:
+                self.writeline("%s.append(" % frame.buffer)
+            else:
+                self.writeline("%s.extend((" % frame.buffer)
+
+            self.indent()
+
+        for item in body:
+            if isinstance(item, list):
+                # A group of constant data to join and output.
+                val = self._output_const_repr(item)
+
+                if frame.buffer is None:
+                    self.writeline("yield " + val)
                 else:
-                    self.writeline('%s.extend((' % frame.buffer)
-                self.indent()
-            for item in body:
-                if isinstance(item, list):
-                    val = repr(concat(item))
-                    if frame.buffer is None:
-                        self.writeline('yield ' + val)
-                    else:
-                        self.writeline(val + ',')
+                    self.writeline(val + ",")
+            else:
+                if frame.buffer is None:
+                    self.writeline("yield ", item)
                 else:
-                    if frame.buffer is None:
-                        self.writeline('yield ', item)
-                    else:
-                        self.newline(item)
-                    close = 1
-                    if frame.eval_ctx.volatile:
-                        self.write('(escape if context.eval_ctx.autoescape'
-                                   ' else to_string)(')
-                    elif frame.eval_ctx.autoescape:
-                        self.write('escape(')
-                    else:
-                        self.write('to_string(')
-                    if self.environment.finalize is not None:
-                        self.write(finalize_src)
-                        close += 1
-                    self.visit(item, frame)
-                    self.write(')' * close)
-                    if frame.buffer is not None:
-                        self.write(',')
-            if frame.buffer is not None:
-                # close the open parentheses
-                self.outdent()
-                self.writeline(len(body) == 1 and ')' or '))')
+                    self.newline(item)
 
-        # otherwise we create a format string as this is faster in that case
-        else:
-            format = []
-            arguments = []
-            for item in body:
-                if isinstance(item, list):
-                    format.append(concat(item).replace('%', '%%'))
-                else:
-                    format.append('%s')
-                    arguments.append(item)
-            self.writeline('yield ')
-            self.write(repr(concat(format)) + ' % (')
-            self.indent()
-            for argument in arguments:
-                self.newline(argument)
-                close = 0
-                if frame.eval_ctx.volatile:
-                    self.write('(escape if context.eval_ctx.autoescape else'
-                               ' to_string)(')
-                    close += 1
-                elif frame.eval_ctx.autoescape:
-                    self.write('escape(')
-                    close += 1
-                if self.environment.finalize is not None:
-                    self.write(finalize_src)
-                    close += 1
-                self.visit(argument, frame)
-                self.write(')' * close + ', ')
+                # A node to be evaluated at runtime.
+                self._output_child_pre(item, frame, finalize)
+                self.visit(item, frame)
+                self._output_child_post(item, frame, finalize)
+
+                if frame.buffer is not None:
+                    self.write(",")
+
+        if frame.buffer is not None:
             self.outdent()
-            self.writeline(')')
+            self.writeline(")" if len(body) == 1 else "))")
 
-        if outdent_later:
+        if frame.require_output_check:
             self.outdent()
 
     def visit_Assign(self, node, frame):
diff --git a/jinja2/nativetypes.py b/jinja2/nativetypes.py
index 46bc568..7128a33 100644
--- a/jinja2/nativetypes.py
+++ b/jinja2/nativetypes.py
@@ -6,7 +6,6 @@ from jinja2 import nodes
 from jinja2._compat import text_type
 from jinja2.compiler import CodeGenerator, has_safe_repr
 from jinja2.environment import Environment, Template
-from jinja2.utils import concat, escape
 
 
 def native_concat(nodes, preserve_quotes=True):
@@ -49,167 +48,36 @@ def native_concat(nodes, preserve_quotes=True):
 
 
 class NativeCodeGenerator(CodeGenerator):
-    """A code generator which avoids injecting ``to_string()`` calls around the
-    internal code Jinja uses to render templates.
+    """A code generator which renders Python types by not adding
+    ``to_string()`` around output nodes, and using :func:`native_concat`
+    to convert complex strings back to Python types if possible.
     """
 
-    def visit_Output(self, node, frame):
-        """Same as :meth:`CodeGenerator.visit_Output`, but do not call
-        ``to_string`` on output nodes in generated code.
-        """
-        if self.has_known_extends and frame.require_output_check:
-            return
-
-        finalize = self.environment.finalize
-        finalize_context = getattr(finalize, 'contextfunction', False)
-        finalize_eval = getattr(finalize, 'evalcontextfunction', False)
-        finalize_env = getattr(finalize, 'environmentfunction', False)
-
-        if finalize is not None:
-            if finalize_context or finalize_eval:
-                const_finalize = None
-            elif finalize_env:
-                def const_finalize(x):
-                    return finalize(self.environment, x)
-            else:
-                const_finalize = finalize
-        else:
-            def const_finalize(x):
-                return x
-
-        # If we are inside a frame that requires output checking, we do so.
-        outdent_later = False
-
-        if frame.require_output_check:
-            self.writeline('if parent_template is None:')
-            self.indent()
-            outdent_later = True
-
-        # Try to evaluate as many chunks as possible into a static string at
-        # compile time.
-        body = []
-
-        for child in node.nodes:
-            try:
-                if const_finalize is None:
-                    raise nodes.Impossible()
-
-                const = child.as_const(frame.eval_ctx)
-                if not has_safe_repr(const):
-                    raise nodes.Impossible()
-            except nodes.Impossible:
-                body.append(child)
-                continue
-
-            # the frame can't be volatile here, because otherwise the as_const
-            # function would raise an Impossible exception at that point
-            try:
-                if frame.eval_ctx.autoescape:
-                    if hasattr(const, '__html__'):
-                        const = const.__html__()
-                    else:
-                        const = escape(const)
-
-                const = const_finalize(const)
-            except Exception:
-                # if something goes wrong here we evaluate the node at runtime
-                # for easier debugging
-                body.append(child)
-                continue
-
-            if body and isinstance(body[-1], list):
-                body[-1].append(const)
-            else:
-                body.append([const])
-
-        # if we have less than 3 nodes or a buffer we yield or extend/append
-        if len(body) < 3 or frame.buffer is not None:
-            if frame.buffer is not None:
-                # for one item we append, for more we extend
-                if len(body) == 1:
-                    self.writeline('%s.append(' % frame.buffer)
-                else:
-                    self.writeline('%s.extend((' % frame.buffer)
-
-                self.indent()
-
-            for item in body:
-                if isinstance(item, list):
-                    val = repr(native_concat(item))
-
-                    if frame.buffer is None:
-                        self.writeline('yield ' + val)
-                    else:
-                        self.writeline(val + ',')
-                else:
-                    if frame.buffer is None:
-                        self.writeline('yield ', item)
-                    else:
-                        self.newline(item)
-
-                    close = 0
-
-                    if finalize is not None:
-                        self.write('environment.finalize(')
-
-                        if finalize_context:
-                            self.write('context, ')
-
-                        close += 1
-
-                    self.visit(item, frame)
-
-                    if close > 0:
-                        self.write(')' * close)
-
-                    if frame.buffer is not None:
-                        self.write(',')
-
-            if frame.buffer is not None:
-                # close the open parentheses
-                self.outdent()
-                self.writeline(len(body) == 1 and ')' or '))')
-
-        # otherwise we create a format string as this is faster in that case
-        else:
-            format = []
-            arguments = []
-
-            for item in body:
-                if isinstance(item, list):
-                    format.append(native_concat(item).replace('%', '%%'))
-                else:
-                    format.append('%s')
-                    arguments.append(item)
-
-            self.writeline('yield ')
-            self.write(repr(concat(format)) + ' % (')
-            self.indent()
-
-            for argument in arguments:
-                self.newline(argument)
-                close = 0
-
-                if finalize is not None:
-                    self.write('environment.finalize(')
-
-                    if finalize_context:
-                        self.write('context, ')
-                    elif finalize_eval:
-                        self.write('context.eval_ctx, ')
-                    elif finalize_env:
-                        self.write('environment, ')
-
-                    close += 1
-
-                self.visit(argument, frame)
-                self.write(')' * close + ', ')
-
-            self.outdent()
-            self.writeline(')')
-
-        if outdent_later:
-            self.outdent()
+    @staticmethod
+    def _default_finalize(value):
+        return value
+
+    def _output_const_repr(self, group):
+        return repr(native_concat(group))
+
+    def _output_child_to_const(self, node, frame, finalize):
+        const = node.as_const(frame.eval_ctx)
+
+        if not has_safe_repr(const):
+            raise nodes.Impossible()
+
+        if isinstance(node, nodes.TemplateData):
+            return const
+
+        return finalize.const(const)
+
+    def _output_child_pre(self, node, frame, finalize):
+        if finalize.src is not None:
+            self.write(finalize.src)
+
+    def _output_child_post(self, node, frame, finalize):
+        if finalize.src is not None:
+            self.write(")")
 
 
 class NativeEnvironment(Environment):
diff --git a/jinja2/runtime.py b/jinja2/runtime.py
index e3aa1f8..cb0bcf5 100644
--- a/jinja2/runtime.py
+++ b/jinja2/runtime.py
@@ -505,7 +505,6 @@ class LoopContext:
     def __iter__(self):
         return self
 
-    @internalcode
     def __next__(self):
         if self._after is not missing:
             rv = self._after
@@ -518,6 +517,7 @@ class LoopContext:
         self._current = rv
         return rv, self
 
+    @internalcode
     def __call__(self, iterable):
         """When iterating over nested data, render the body of the loop
         recursively with the given inner iterable data.
diff --git a/tests/test_api.py b/tests/test_api.py
index ec93db2..2eb9ce7 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -11,6 +11,7 @@
 import os
 import tempfile
 import shutil
+from io import StringIO
 
 import pytest
 from jinja2 import Environment, Undefined, ChainableUndefined, \
@@ -206,25 +207,24 @@ class TestMeta(object):
 @pytest.mark.api
 @pytest.mark.streaming
 class TestStreaming(object):
-
     def test_basic_streaming(self, env):
-        tmpl = env.from_string("<ul>{% for item in seq %}<li>{{ loop.index "
-                               "}} - {{ item }}</li>{%- endfor %}</ul>")
-        stream = tmpl.stream(seq=list(range(4)))
-        assert next(stream) == '<ul>'
-        assert next(stream) == '<li>1 - 0</li>'
-        assert next(stream) == '<li>2 - 1</li>'
-        assert next(stream) == '<li>3 - 2</li>'
-        assert next(stream) == '<li>4 - 3</li>'
-        assert next(stream) == '</ul>'
+        t = env.from_string(
+            "<ul>{% for item in seq %}<li>{{ loop.index }} - {{ item }}</li>"
+            "{%- endfor %}</ul>"
+        )
+        stream = t.stream(seq=list(range(3)))
+        assert next(stream) == "<ul>"
+        assert "".join(stream) == "<li>1 - 0</li><li>2 - 1</li><li>3 - 2</li></ul>"
 
     def test_buffered_streaming(self, env):
-        tmpl = env.from_string("<ul>{% for item in seq %}<li>{{ loop.index "
-                               "}} - {{ item }}</li>{%- endfor %}</ul>")
-        stream = tmpl.stream(seq=list(range(4)))
+        tmpl = env.from_string(
+            "<ul>{% for item in seq %}<li>{{ loop.index }} - {{ item }}</li>"
+            "{%- endfor %}</ul>"
+        )
+        stream = tmpl.stream(seq=list(range(3)))
         stream.enable_buffering(size=3)
-        assert next(stream) == u'<ul><li>1 - 0</li><li>2 - 1</li>'
-        assert next(stream) == u'<li>3 - 2</li><li>4 - 3</li></ul>'
+        assert next(stream) == u'<ul><li>1'
+        assert next(stream) == u' - 0</li>'
 
     def test_streaming_behavior(self, env):
         tmpl = env.from_string("")
diff --git a/tests/test_async.py b/tests/test_async.py
index 5f331a5..b71f094 100644
--- a/tests/test_async.py
+++ b/tests/test_async.py
@@ -102,13 +102,12 @@ def test_async_iteration_in_templates():
 
 
 def test_async_iteration_in_templates_extended():
-    t = Template('{% for x in rng %}{{ loop.index0 }}/{{ x }}{% endfor %}',
-                 enable_async=True)
-    async def async_iterator():
-        for item in [1, 2, 3]:
-            yield item
-    rv = list(t.generate(rng=async_iterator()))
-    assert rv == ['0/1', '1/2', '2/3']
+    t = Template(
+        "{% for x in rng %}{{ loop.index0 }}/{{ x }}{% endfor %}", enable_async=True
+    )
+    stream = t.generate(rng=auto_aiter(range(1, 4)))
+    assert next(stream) == "0"
+    assert "".join(stream) == "/11/22/3"
 
 
 @pytest.fixture