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 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_preserves_annotations(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
anno.setanno(node.body[0], 'foo', 'bar')
anno.setanno(node.body[0], 'baz', 1)
new_node = ast_util.copy_clean(node, preserve_annos={'foo'})
self.assertEqual(anno.getanno(new_node.body[0], 'foo'), 'bar')
self.assertFalse(anno.hasanno(new_node.body[0], 'baz'))
def test_copy_clean(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
setattr(node.body[0], '__foo', 'bar')
new_node = ast_util.copy_clean(node)
self.assertIsNot(new_node, node)
self.assertIsNot(new_node.body[0], node.body[0])
self.assertFalse(hasattr(new_node.body[0], '__foo'))
def test_copy_clean_preserves_annotations(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
anno.setanno(node.body[0], 'foo', 'bar')
anno.setanno(node.body[0], 'baz', 1)
new_node = ast_util.copy_clean(node, preserve_annos={'foo'})
self.assertEqual(anno.getanno(new_node.body[0], 'foo'), 'bar')
self.assertFalse(anno.hasanno(new_node.body[0], 'baz'))
def test_copy_clean(self):
node = parser.parse_str(
textwrap.dedent("""
def f(a):
return a + 1
"""))
setattr(node.body[0], '__foo', 'bar')
new_node = ast_util.copy_clean(node)
self.assertIsNot(new_node, node)
self.assertIsNot(new_node.body[0], node.body[0])
self.assertFalse(hasattr(new_node.body[0], '__foo'))
def _prepare_replacement(self, replaced, key):
"""Prepares a replacement AST that's safe to swap in for a node.
Args:
replaced: ast.AST, the node being replaced
key: Hashable, the key of the replacement AST
Returns:
ast.AST, the replacement AST
"""
repl = self.replacements[key]
new_nodes = ast_util.copy_clean(repl, preserve_annos=self.preserved_annos)
if isinstance(new_nodes, gast.AST):
new_nodes = [new_nodes]
return new_nodes
def _prepare_replacement(self, replaced, key):
"""Prepares a replacement AST that's safe to swap in for a node.
Args:
replaced: ast.AST, the node being replaced
key: Hashable, the key of the replacement AST
Returns:
ast.AST, the replacement AST
"""
repl = self.replacements[key]
new_nodes = ast_util.copy_clean(repl, preserve_annos=self.preserved_annos)
if isinstance(new_nodes, gast.AST):
new_nodes = [new_nodes]
return new_nodes
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