unconditionally override lineno and col_offset on generated ast

This commit is contained in:
Benjamin Peterson 2011-05-19 21:49:37 -05:00
parent bf039fea74
commit e0c128beec
1 changed files with 14 additions and 3 deletions

View File

@ -67,6 +67,19 @@ binop_map = {
} }
def set_location(node, lineno, col_offset):
"""Set node location information recursively."""
def _fix(node, lineno, col_offset):
if "lineno" in node._attributes:
node.lineno = lineno
if "col_offset" in node._attributes:
node.col_offset = col_offset
for child in ast.iter_child_nodes(node):
_fix(child, lineno, col_offset)
_fix(node, lineno, col_offset)
return node
class AssertionRewriter(ast.NodeVisitor): class AssertionRewriter(ast.NodeVisitor):
def run(self, mod): def run(self, mod):
@ -196,9 +209,7 @@ class AssertionRewriter(ast.NodeVisitor):
self.statements.append(delete) self.statements.append(delete)
# Fix line numbers. # Fix line numbers.
for stmt in self.statements: for stmt in self.statements:
stmt.lineno = assert_.lineno set_location(stmt, assert_.lineno, assert_.col_offset)
stmt.col_offset = assert_.col_offset
ast.fix_missing_locations(stmt)
return self.statements return self.statements
def visit_Name(self, name): def visit_Name(self, name):