def visit_nodelist(self, nodelist):
for i in range(len(nodelist)):
node = nodelist[i]
if isinstance(node, gast.If):
true_branch_returns = isinstance(node.body[-1], gast.Return)
false_branch_returns = len(node.orelse) and isinstance(
node.orelse[-1], gast.Return)
# If the last node in the if body is a return,
# then every line after this if statement effectively
# belongs in the else.
if true_branch_returns and not false_branch_returns:
for j in range(i + 1, len(nodelist)):
nodelist[i].orelse.append(
ast_util.copy_clean(nodelist[j]))
if nodelist[i + 1:]:
self.changes_made = True
return nodelist[:i + 1]
elif not true_branch_returns and false_branch_returns:
for j in range(i + 1, len(nodelist)):
nodelist[i].body.append(
ast_util.copy_clean(nodelist[j]))
if nodelist[i + 1:]:
self.changes_made = True
return nodelist[:i + 1]
elif true_branch_returns and false_branch_returns:
if nodelist[i + 1:]:
raise ValueError(
'Unreachable code after conditional where both branches return.'
)
return nodelist
elif isinstance(node, gast.Return) and nodelist[i + 1:]:
raise ValueError(
'Cannot have statements after a return in the same basic block'
)
return nodelist
def test_copy_clean(self):
ret = ast.Return(
ast.BinOp(op=ast.Add(),
left=ast.Name(id='a', ctx=ast.Load()),
right=ast.Num(1)))
setattr(ret, '__foo', 'bar')
node = ast.FunctionDef(name='f',
args=ast.arguments(
args=[ast.Name(id='a', ctx=ast.Param())],
vararg=None,
kwarg=None,
defaults=[]),
body=[ret],
decorator_list=[],
returns=None)
new_node = ast_util.copy_clean(node)
self.assertFalse(node is new_node)
self.assertFalse(ret is new_node.body[0])
self.assertFalse(hasattr(new_node.body[0], '__foo'))
def test_copy_clean(self):
ret = ast.Return(
ast.BinOp(
op=ast.Add(),
left=ast.Name(id='a', ctx=ast.Load()),
right=ast.Num(1)))
setattr(ret, '__foo', 'bar')
node = ast.FunctionDef(
name='f',
args=ast.arguments(
args=[ast.Name(id='a', ctx=ast.Param())],
vararg=None,
kwarg=None,
defaults=[]),
body=[ret],
decorator_list=[],
returns=None)
new_node = ast_util.copy_clean(node)
self.assertFalse(node is new_node)
self.assertFalse(ret is new_node.body[0])
self.assertFalse(hasattr(new_node.body[0], '__foo'))
def visit_Name(self, node):
if node.id not in self.replacements:
return node
new_nodes = ast_util.copy_clean(self.replacements[node.id])
if isinstance(new_nodes, gast.AST):
new_nodes = [new_nodes]
# Preserve the target context.
for n in new_nodes:
if isinstance(n, gast.Tuple):
for e in n.elts:
self._set_inner_child_context(e, node.ctx)
if isinstance(n, gast.Attribute):
# For attributes, the inner Name node receives the context, while the
# outer ones have it set to Load.
self._set_inner_child_context(n, node.ctx)
else:
n.ctx = node.ctx
if len(new_nodes) == 1:
new_nodes, = new_nodes
return new_nodes
def visit_Name(self, node):
if node.id not in self.replacements:
return node
new_nodes = ast_util.copy_clean(self.replacements[node.id])
if isinstance(new_nodes, gast.AST):
new_nodes = [new_nodes]
# Preserve the target context.
for n in new_nodes:
if isinstance(n, gast.Tuple):
for e in n.elts:
self._set_inner_child_context(e, node.ctx)
if isinstance(n, gast.Attribute):
# For attributes, the inner Name node receives the context, while the
# outer ones have it set to Load.
self._set_inner_child_context(n, node.ctx)
else:
n.ctx = node.ctx
if len(new_nodes) == 1:
new_nodes, = new_nodes
return new_nodes