def visit_If(self, node):
if not coloring.is_green(node):
return node
node.body = self._transform_sequence(node.body)
node.orelse = self._transform_sequence(node.orelse)
# transform node test, e.g. if it's a call, this makes sure it's lifted
node.test = self.visit(node.test)
# lift the test (possibly again). We simply always do this,
# because we can't know whether the value in it is a bool or
# not.
#
# TODO: we can be a lot smarter when there are
# Compare(...constant..) exprs, and transform to SFN's string
# equality, numeric equality etc operators
testVar = gensym.gensym("test")
lifted = ast.Assign(targets=[ast.Name(id=testVar, ctx=ast.Store())],
value=ast.Call(func=ast.Name(id='bool',
ctx=ast.Load()),
args=[node.test],
keywords=[]))
ast.copy_location(lifted, node.test)
self.pre_statements.append(lifted)
node.test = ast.copy_location(ast.Name(id=testVar, ctx=ast.Load()),
node)
return node
def visit_Try(self, node):
if not coloring.is_green(node):
return
node.body = self._transform_sequence(node.body)
for h in node.handlers:
h.body = self._transform_sequence(h.body)
return node
def transform_node(node: ast.AST) -> CAST:
if not coloring.is_green(node):
return PythonAST([node]).set_loc(node)
else:
#print("Transforming a %s" % node.__class__.__name__)
n = transform_green_node(node)
if not isinstance(node, ast.Module):
n.set_loc(node)
return n
def visit_Expr(self, node):
# Special-case void-context calls. This avoids a whole
# unnecessary lambda for the very common case of invoking a
# task and discarding the result.
v = node.value
if isinstance(v, ast.Call):
if not coloring.is_green(v):
return
self._lift_call_args(v)
return node
def visit_Return(self, node):
# blue returns are never transformed;
# green returns are transformed UNLESS:
# (a) they return a constant primitive value (num, string, bool, nothing)
# (b) they return a variable reference (with no other expression)
if not coloring.is_green(node):
return
v = node.value
if v == None:
return node
if isinstance(v, ast.Num) or isinstance(v, ast.Str) or isinstance(
v, ast.NameConstant):
return node
if isinstance(v, ast.Name) and isinstance(v.ctx, ast.Load):
return node
# return <expr> becomes:
# ...
# env[var] = transform[<expr>]
# return env[var]
#
tmpVar = gensym.gensym("ret")
# node.value may contain green nodes, so we need to visit it
retvalue = self.visit(node.value)
if isinstance(retvalue, ast.Name) and isinstance(
retvalue.ctx, ast.Load):
# no further transform needed
replacement = ast.copy_location(ast.Return(value=retvalue), node)
replacement.color = True
return replacement
# we got something more than a name (some sort of expression).
# lift it out to an assignment.
lifted = ast.Assign(targets=[ast.Name(id=tmpVar, ctx=ast.Store())],
value=retvalue)
ast.copy_location(lifted, node)
# I think we don't need to color this because whatever green
# is in there has been lifted? Need to prove/disprove this.
self.pre_statements.append(lifted)
replacement = ast.Return(value=ast.Name(id=tmpVar, ctx=ast.Load()))
ast.copy_location(replacement, node)
replacement.color = True
return replacement
def visit_Assign(self, node):
# assignment transforms:
# if blue do nothing
# if green and rhs is Call, do nothing
# if green and rhs is anything else (e.g. a binop), change
# color to blue: the call will be lifted out of this assign,
# whatever expr it's in.
if not coloring.is_green(node):
return node
if isinstance(node.value, ast.Call):
return node
node.color = False
return self.generic_visit(node)
def visit_Call(self, node):
"""Transform a call into (a) assignments for args (b) standalone call
(variables in, out) (c) replacement node as tmp var ref
"""
if not coloring.is_green(node):
return node
self._lift_call_args(node)
# now, the call itself
resultVar = gensym.gensym("call")
lifted = ast.Assign(targets=[ast.Name(id=resultVar, ctx=ast.Store())],
value=node)
ast.copy_location(lifted, node)
lifted.color = True
self.pre_statements.append(lifted)
replacement = ast.copy_location(ast.Name(id=resultVar, ctx=ast.Load()),
node)
return replacement
def transform_list(nodes: typing.List[ast.AST]) -> typing.List[CAST]:
if nodes == None:
return None
result: typing.List[CAST] = []
curr_blue = None
for n in nodes:
#print("--- visiting ----- %s" % n.__class__.__name__)
if coloring.is_green(n):
# add curr_blue list to ast
if curr_blue != None:
pa = PythonAST(curr_blue)
result.append(pa)
curr_blue = None
# handle green node
result.append(transform_node(n))
else:
if curr_blue == None:
curr_blue = []
curr_blue.append(n)
if curr_blue: # last node was blue
result.append(PythonAST(curr_blue))
return result
def visit_While(self, node):
if not coloring.is_green(node):
return node
# testVar = not bool(...test)
testVar = gensym.gensym("test")
testStatement = ast.Assign(
targets=[ast.Name(id=testVar, ctx=ast.Store())],
value=ast.UnaryOp(op=ast.Not(),
operand=ast.Call(func=ast.Name(id='bool',
ctx=ast.Load()),
args=[node.test],
keywords=[])))
ast.copy_location(testStatement, node)
# TODO testStatement needs to be colored appropriately
# based on whether test contains green calls
# "if <testVar>: break"
breakNode = ast.Break()
ast.copy_location(breakNode, node)
breakNode.color = True
loopIf = ast.If(test=ast.Name(id=testVar, ctx=ast.Load()),
body=[breakNode],
orelse=[])
ast.copy_location(loopIf, node)
loopIf.color = True
# build the new body
newbody = []
newbody.append(testStatement)
newbody.append(loopIf)
newbody.extend(node.body)
node.body = self._transform_sequence(newbody)
node.test = ast.copy_location(ast.NameConstant(value=True), node)
return node
def visit_For(self, node):
if not coloring.is_green(node):
return node
node.body = self._transform_sequence(node.body)
node.iter = self.generic_visit(node.iter)
return node
def visit_FunctionDef(self, node):
if not coloring.is_green(node):
return node
node.body = self._transform_sequence(node.body)
return node
def visit_FunctionDef(self, node):
oldcolor = self.green
self.green = coloring.is_green(node)
self.generic_visit(node)
self.green = oldcolor
return node
def transform_green_node(node: ast.AST) -> CAST:
if not coloring.is_green(node):
raise Exception("Transform error")
c = node.__class__
if c == ast.FunctionDef:
return Def(name = node.name,
args = [Arg(a.arg) for a in node.args.args],
body = transform_list(node.body))
elif c == ast.For:
return ForLoop(target = astor.to_source(node.target).rstrip(),
collection = astor.to_source(node.iter).rstrip(),
body = transform_list(node.body))
elif c == ast.While:
return WhileLoop(test = node.test,
body = transform_list(node.body))
elif c == ast.If:
return If(test = node.test,
thenBody = transform_list(node.body),
elseBody = transform_list(node.orelse))
elif c == ast.Call:
argVars = []
# args should only be vars -- lift.calls lifts out all expressions in (green) function arguments
# TODO: this pessimises inline constants by adding unnecessary assignments
for a in node.args:
argVars.append(varFromEnv(a))
result = Call(func = astor.to_source(node.func).rstrip(),
args = argVars)
keywordArgs = []
for k in node.keywords:
arg = k.arg.lower()
if arg == 'timeout' or arg == 'timeoutseconds':
result.timeoutSec = getNumber(k.value)
elif arg == 'heartbeat' or arg == 'heartbeatseconds':
result.heartbeatSec = getNumber(k.value)
elif arg == 'retry':
result.retry = getRetrier(k.value)
return result
elif c == ast.Module:
return Module(body = transform_list(node.body))
elif c == ast.Assign:
# TODO for now just support single assignment
return Assign(varFromEnv(node.targets[0]), transform_node(node.value))
elif c == ast.Expr:
# I'm not sure why this Expr node exists and not just directly
# the Call or whatever. So we don't model it in the CAST, we
# just pass thru to whatever node is inside it (e.g. Call)
return transform_node(node.value)
elif c == ast.Return:
return Return(node.value)
elif c == ast.Break:
return Break()
elif c == ast.Try:
handlers = []
for eh in node.handlers:
types = []
if isinstance(eh.type, ast.Tuple):
types += [t.id for t in eh.type.elts]
else:
types += [astor.to_source(eh.type).rstrip()]
handlers.append(handler(types = types,
name = eh.name,
body = transform_list(eh.body)))
return Try(body = transform_list(node.body), handlers = handlers)
else:
raise Exception("Unhandled node type: %s" % c.__name__)