def is_atom_string(node):
try:
power = get_power(node)
return (len(power) == 2 and
power[1][0] == symbol.atom and
power[1][1][0] == token.STRING)
except IndexError:
return False
def __init__(self, tree):
self.tree = tree
self.template_type_stack = [None]
self.module_start = statement_tree(
'from qpy.quoted import join_xml as _qpy_join_xml, '
'join_str as _qpy_join_str, xml as _qpy_xml')
self.template_start = statement_tree(
'qpy_accumulation=[];qpy_append=qpy_accumulation.append')
self.return_xml = statement_tree(
'return _qpy_join_xml(qpy_accumulation)')
self.return_str = statement_tree(
'return _qpy_join_str(qpy_accumulation)')
self.qpy_append_expr_stmt = statement_tree('qpy_append(X)')[1][1][1]
assert self.qpy_append_expr_stmt[0] == symbol.expr_stmt
self.xml_power = get_power(expr_tree('_qpy_xml("X")'))
self.template_type_stack = [None]
self.traverse_node(self.tree)
self.line_number = 1
self.rationalize_line_numbers(self.tree)
def traverse_node(self, node):
if not isinstance(node, list):
return
# If this is a funcdef, push 'xml', 'str', or None
if node[0] == symbol.funcdef:
function_name = get_funcdef_function_name_child(node)[1]
if function_name.endswith('__xml_template__'):
self.template_type_stack.append('xml')
elif function_name.endswith('__str_template__'):
self.template_type_stack.append('str')
else:
self.template_type_stack.append(None)
# Traverse down before doing modifications.
for child in node[1:]:
self.traverse_node(child)
# Modify the node as necessary.
if node[0] == symbol.file_input:
# Insert module-level import statement.
# Skip over the module docstring and any __future__ imports.
for index, child in enumerate(node):
if index == 0:
continue
if is_atom_string(child):
continue
if is_future_import_statement(child):
continue
node.insert(index, deepcopy(self.module_start))
break
elif self.template_type_stack[-1] is None:
pass # We're not in a template, so we're done.
elif node[0] == symbol.expr_stmt and len(node) == 2:
# Wrap this expression statement in a qpy_append call.
stmt = deepcopy(self.qpy_append_expr_stmt)
argument = get_argument(get_power(stmt))
assert argument[1][0] == node[1][1][0]
argument[1] = node[1][1]
assert node[1][0] == stmt[1][0]
node[1] = stmt[1]
elif node[0] == symbol.funcdef:
# This is a new template.
# Insert the initialization of qpy_accumulation and qpy_append.
func_suite = node[-1]
for j, child in enumerate(func_suite[1:]):
if child[0] == symbol.stmt:
func_suite.insert(j+1, deepcopy(self.template_start))
break
# Add the appropriate return statement and patch function name.
function_name_node = get_funcdef_function_name_child(node)
function_name = function_name_node[1]
if self.template_type_stack[-1] == 'xml':
return_accumulation = deepcopy(self.return_xml)
# trim __xml_template__
function_name_node[1] = function_name_node[1][:-16]
else:
assert self.template_type_stack[-1] == 'str'
return_accumulation = deepcopy(self.return_str)
# trim __str_template__
function_name_node[1] = function_name_node[1][:-16]
func_suite.insert(-1, return_accumulation)
elif (self.template_type_stack[-1] == 'xml' and
node[0] == symbol.power and
node[1][0] == symbol.atom and
node[1][1][0] == token.STRING):
# node looks like
# [power [atom [STRING "Z" 1]] ...]
xml_power = deepcopy(self.xml_power)
# xml_power looks like
# [power
# [atom [STRING "_qpy_xml" 1]
# [trailer
# LPAR ...
# [arglist ... [power [atom [STRING "X"]]]]
# RPAR]]]
argument = get_argument(xml_power)
argument_power = get_power(argument)
assert argument_power[1][0] == symbol.atom
assert argument_power[1][0] == node[1][0]
argument_power[1] = deepcopy(node[1]) # replace "X"
node[1] = xml_power[1]
node.insert(2, xml_power[2])
# Pop the stack.
if node[0] == symbol.funcdef:
self.template_type_stack.pop()