class QueryDirectiveRewriter(P.NodeTransformer):
"""Take QUERY(<source>, **<kargs>) directives and replace the
first argument with the corresponding parsed Python AST.
"""
# We'll use NodeTransformer rather than MacroProcessor because
# MacroProcessor makes it hard to reconstruct the node.
pat = P.Expr(
P.Call(P.Name('QUERY', P.Load()), [P.Str(P.PatVar('_source'))],
P.PatVar('_keywords'), None, None))
def visit_Expr(self, node):
match = P.match(self.pat, node)
if match is not None:
query = P.Parser.pe(node.value.args[0].s)
call = node.value._replace(args=[query])
node = node._replace(value=call)
return node
def postprocess_var_decls(tree, decls):
"""Prepend global variable declarations to the program. Each given
declaration is a triple of a variable name, type name (e.g., 'Set'
or 'Map'), and a declaration comment string.
"""
assert isinstance(tree, P.Module)
header = ()
for var_name, type_name, comment in decls:
header += P.Parser.pc('''
COMMENT(_S)
_VAR = _TYPE()
''',
subst={
'_S': P.Str(comment),
'_VAR': var_name,
'_TYPE': type_name
})
tree = tree._replace(body=header + tree.body)
return tree
def postprocess_header(tree, header):
"""Add comment lines for each string in header."""
header = tuple(
P.Parser.ps('COMMENT(_S)', subst={'_S': P.Str(line)})
for line in header)
return tree._replace(body=header + tree.body)
