google / pasta

Library to refactor python code through AST manipulation.
Apache License 2.0
341 stars 40 forks source link

IndentationError when appending new node to function body #82

Closed cript0nauta closed 4 years ago

cript0nauta commented 4 years ago

Hi! I found a bug in pasta when I try to add a node to a function body (I think it's the same with every node class that has a body). I built the following script to reproduce the bug:

import ast
import pasta

source_code = """
if 1:  # This is required
    def f():
        a=1
"""

class SampleTransformer(ast.NodeTransformer):
    def visit_FunctionDef(self, node):
        return_node = ast.copy_location(
            ast.Return(value=None), node.body[0])
        # return_node = ast.Return(value=None)  # This produces the same error than above
        node.body.append(return_node)
        return node

tree = pasta.parse(source_code)
SampleTransformer().visit(tree)

print(pasta.dump(tree))

The expected output of it should be:

if 1:  # This is required
    def f():
        a=1
        return

But instead it produces a syntactically invalid code:

if 1:  # This is required
    def f():
        a=1
            return
soupytwist commented 4 years ago

That was an ugly bug which I'm surprised didn't show sooner. Thanks for the report!

soupytwist commented 4 years ago

And FYI-- ast.copy_location does nothing here and isn't needed.

cript0nauta commented 4 years ago

Thanks for the fast response, it's working OK now!