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()