- improve the line counting to take place in the printer; the state

is also tracked for the template as a whole whereas before it was
only occurring for the top level node and was therefore incomplete

2 files changed, 42 insertions, 36 deletions

diff --git a/mako/codegen.py b/mako/codegen.py
index ec587ba..6380b40 100644
--- a/mako/codegen.py
+++ b/mako/codegen.py
@@ -99,11 +99,8 @@ class _GenerateRenderMethod(object):
     """
     def __init__(self, printer, compiler, node):
         self.printer = printer
-        self.last_source_line = -1
         self.compiler = compiler
         self.node = node
-        self.source_map = {}
-        self.boilerplate_map = []
         self.identifier_stack = [None]
         self.in_def = isinstance(node, (parsetree.DefTag, parsetree.BlockTag))
 
@@ -152,15 +149,14 @@ class _GenerateRenderMethod(object):
             self.write_metadata_struct()
 
     def write_metadata_struct(self):
-        self.source_map[self.printer.lineno] = self.last_source_line
+        self.printer.source_map[self.printer.lineno] = self.printer.last_source_line
         struct = {
             "filename": self.compiler.filename,
             "uri": self.compiler.uri,
             "source_encoding": self.compiler.source_encoding,
-            "line_map": self.source_map,
-            "boilerplate_lines": self.boilerplate_map
+            "line_map": self.printer.source_map,
+            "boilerplate_lines": self.printer.boilerplate_map
         }
-        self.mark_boilerplate()
         self.printer.writelines(
             '"""',
             '__M_BEGIN_METADATA',
@@ -274,7 +270,6 @@ class _GenerateRenderMethod(object):
 
         this could be the main render() method or that of a top-level def."""
 
-        self.mark_boilerplate()
         if self.in_def:
             decorator = node.decorator
             if decorator:
@@ -323,13 +318,12 @@ class _GenerateRenderMethod(object):
         """write module-level template code, i.e. that which
         is enclosed in <%! %> tags in the template."""
         for n in module_code:
-            self.write_source_comment(n)
+            self.printer.start_source(n.lineno)
             self.printer.write_indented_block(n.text)
 
     def write_inherit(self, node):
         """write the module-level inheritance-determination callable."""
 
-        self.mark_boilerplate()
         self.printer.writelines(
             "def _mako_inherit(template, context):",
                 "_mako_generate_namespaces(context)",
@@ -340,7 +334,6 @@ class _GenerateRenderMethod(object):
 
     def write_namespaces(self, namespaces):
         """write the module-level namespace-generating callable."""
-        self.mark_boilerplate()
         self.printer.writelines(
             "def _mako_get_namespace(context, name):",
                 "try:",
@@ -356,7 +349,7 @@ class _GenerateRenderMethod(object):
         for node in namespaces.values():
             if 'import' in node.attributes:
                 self.compiler.has_ns_imports = True
-            self.write_source_comment(node)
+            self.printer.start_source(node.lineno)
             if len(node.nodes):
                 self.printer.writeline("def make_namespace():")
                 export = []
@@ -558,16 +551,6 @@ class _GenerateRenderMethod(object):
 
         self.printer.writeline("__M_writer = context.writer()")
 
-    def write_source_comment(self, node):
-        """write a source comment containing the line number of the
-        corresponding template line."""
-        if self.last_source_line != node.lineno:
-            self.source_map[self.printer.lineno] = node.lineno
-            self.last_source_line = node.lineno
-
-    def mark_boilerplate(self):
-        self.boilerplate_map.append(self.printer.lineno)
-
     def write_def_decl(self, node, identifiers):
         """write a locally-available callable referencing a top-level def"""
         funcname = node.funcname
@@ -635,7 +618,6 @@ class _GenerateRenderMethod(object):
         writes code to retrieve captured content, apply filters, send proper
         return value."""
 
-        self.mark_boilerplate()
         if not buffered and not cached and not filtered:
             self.printer.writeline("return ''")
             if callstack:
@@ -786,7 +768,7 @@ class _GenerateRenderMethod(object):
         return target
 
     def visitExpression(self, node):
-        self.write_source_comment(node)
+        self.printer.start_source(node.lineno)
         if len(node.escapes) or \
                 (
                     self.compiler.pagetag is not None and
@@ -808,7 +790,7 @@ class _GenerateRenderMethod(object):
                 self.printer.writeline("loop = __M_loop._exit()")
                 self.printer.writeline(None)
         else:
-            self.write_source_comment(node)
+            self.printer.start_source(node.lineno)
             if self.compiler.enable_loop and node.keyword == 'for':
                 text = mangle_mako_loop(node, self.printer)
             else:
@@ -830,7 +812,7 @@ class _GenerateRenderMethod(object):
                 self.printer.writeline("pass")
 
     def visitText(self, node):
-        self.write_source_comment(node)
+        self.printer.start_source(node.lineno)
         self.printer.writeline("__M_writer(%s)" % repr(node.content))
 
     def visitTextTag(self, node):
@@ -856,7 +838,7 @@ class _GenerateRenderMethod(object):
 
     def visitCode(self, node):
         if not node.ismodule:
-            self.write_source_comment(node)
+            self.printer.start_source(node.lineno)
             self.printer.write_indented_block(node.text)
 
             if not self.in_def and len(self.identifiers.locally_assigned) > 0:
@@ -873,7 +855,7 @@ class _GenerateRenderMethod(object):
                       ','.join([repr(x) for x in node.declared_identifiers()]))
 
     def visitIncludeTag(self, node):
-        self.write_source_comment(node)
+        self.printer.start_source(node.lineno)
         args = node.attributes.get('args')
         if args:
             self.printer.writeline(
@@ -891,7 +873,6 @@ class _GenerateRenderMethod(object):
         pass
 
     def visitBlockTag(self, node):
-        self.mark_boilerplate()
         if node.is_anonymous:
             self.printer.writeline("%s()" % node.funcname)
         else:
@@ -961,7 +942,6 @@ class _GenerateRenderMethod(object):
             n.accept_visitor(self)
         self.identifier_stack.pop()
 
-        self.mark_boilerplate()
         self.write_def_finish(node, buffered, False, False, callstack=False)
         self.printer.writelines(
             None,
@@ -975,7 +955,7 @@ class _GenerateRenderMethod(object):
                 "runtime.Namespace('caller', context, "
                                   "callables=ccall(__M_caller))",
             "try:")
-        self.write_source_comment(node)
+        self.printer.start_source(node.lineno)
         self.printer.writelines(
                 "__M_writer(%s)" % self.create_filter_callable(
                                                     [], node.expression, True),
diff --git a/mako/pygen.py b/mako/pygen.py
index d6559e5..62967e3 100644
--- a/mako/pygen.py
+++ b/mako/pygen.py
@@ -37,10 +37,36 @@ class PythonPrinter(object):
 
         self._reset_multi_line_flags()
 
-    def write_blanks(self, num=1):
-        self.stream.write("\n" * num)
+        # marker for template source lines; this
+        # is part of source/template line mapping
+        self.last_source_line = -1
+
+        self.last_boilerplate_line = -1
+
+        # mapping of generated python lines to template
+        # source lines
+        self.source_map = {}
+
+        # list of "boilerplate" lines, these are lines
+        # that precede/follow a set of template source-mapped lines
+        self.boilerplate_map = []
+
+
+    def _update_lineno(self, num):
+        if self.last_boilerplate_line <= self.last_source_line:
+            self.boilerplate_map.append(self.lineno)
+            self.last_boilerplate_line = self.lineno
         self.lineno += num
 
+    def start_source(self, lineno):
+        if self.last_source_line != lineno:
+            self.source_map[self.lineno] = lineno
+            self.last_source_line = lineno
+
+    def write_blanks(self, num):
+        self.stream.write("\n" * num)
+        self._update_lineno(num)
+
     def write_indented_block(self, block):
         """print a line or lines of python which already contain indentation.
 
@@ -98,7 +124,7 @@ class PythonPrinter(object):
 
         # write the line
         self.stream.write(self._indent_line(line) + "\n")
-        self.lineno += 1
+        self._update_lineno(1)
 
         # see if this line should increase the indentation level.
         # note that a line can both decrase (before printing) and
@@ -218,13 +244,13 @@ class PythonPrinter(object):
         for entry in self.line_buffer:
             if self._in_multi_line(entry):
                 self.stream.write(entry + "\n")
-                self.lineno += 1
+                self._update_lineno(1)
             else:
                 entry = entry.expandtabs()
                 if stripspace is None and re.search(r"^[ \t]*[^# \t]", entry):
                     stripspace = re.match(r"^([ \t]*)", entry).group(1)
                 self.stream.write(self._indent_line(entry, stripspace) + "\n")
-                self.lineno += 1
+                self._update_lineno(1)
 
         self.line_buffer = []
         self._reset_multi_line_flags()