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