def preprocess_var_decls(tree):
"""Eliminate global variable declarations of the form
S = Set()
M = Map()
and return a list of pairs of variables names and type names (i.e.,
'Set' or 'Map').
"""
assert isinstance(tree, P.Module)
pat = P.Assign([P.Name(P.PatVar('_VAR'), P.Wildcard())],
P.Call(P.Name(P.PatVar('_KIND'), P.Load()), [], [], None,
None))
decls = OrderedSet()
body = []
changed = False
for stmt in tree.body:
match = P.match(pat, stmt)
if match is not None:
var, kind = match['_VAR'], match['_KIND']
if kind not in ['Set', 'Map']:
raise L.ProgramError(
'Unknown declaration initializer {}'.format(kind))
decls.add((var, kind))
changed = True
else:
body.append(stmt)
if changed:
tree = tree._replace(body=body)
return tree, list(decls)
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
