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(self):
node = parser.parse(
textwrap.dedent("""
def f(a):
return a + 1
"""))
setattr(node, '__foo', 'bar')
new_node = ast_util.copy_clean(node)
self.assertIsNot(new_node, node)
self.assertFalse(hasattr(new_node, '__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_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
