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
def visit_Import(self, node):
pat = P.Import([P.alias('incoq.runtime', P.PatVar('_ALIAS'))])
match = P.match(pat, node)
if match is None:
return node
# Remove the import. Record the alias to remove its uses.
if match['_ALIAS'] is not None:
qual = P.Name(match['_ALIAS'], P.Load())
self.quals.add(qual)
return []
def visit_Attribute(self, node):
# Note: As written, this will remove all aliases, not just the
# prefix alias. That is, if A is an alias for the runtime in
# quals, this will rewrite A.A...A.foo as foo. This shouldn't
# be a problem because we shouldn't see aliases chained in this
# manner.
node = self.generic_visit(node)
# Check prefix against each alias and the fully qualified path.
# If none match, no change.
for qual in self.quals:
pat = P.Attribute(qual, P.PatVar('_ATTR'), P.Load())
match = P.match(pat, node)
if match is not None:
break
else:
return node
return P.Name(match['_ATTR'], P.Load())