Visiting a assert statement within a functiondef using Python ast

Question

I am currently doing some code manipulation and I am unable to visit an "Assert" statement within a "FunctionDef". How I visit an internal Assert node within a Function node?

Here is a sample code I am modifying

assert True == True
def test_code():
    result = 5
    assert result == 5

Here is my NodeTransformer

class TransformMe(ast.NodeTransformer):
    def visit_Assert(self, node):
        return None  # Some processing here

    def visit_FunctionDef(self, node):
        for line_node in node.body:
            if isinstance(line_node, ast.Assert):
                line_node = ast.NodeTransformer.generic_visit(self, line_node)
            for test_node in ast.walk(line_node):
                # More processing here
                pass
        return node

Processing it with

module_node = ast.parse(code)
new_module_node = TransformMe().visit(module_node)
print(ast.unparse(new_module_node))

Current output is

def test_code():
    result = 5
    assert result == 5

But the output I want is

def test_code():
    result = 5

How do I make this work? How do I visit and process an internal Assert node?

Did you intend to write `line_node = self.visit_Assert(line_node)`? — mkrieger1, Jul 02 '23 at 21:00

score 0 · Accepted Answer · answered Jul 02 '23 at 21:40

Thanks to @mkrieger1 's comment, there is the working NodeTransformer.

   class TransformMe(ast.NodeTransformer):
    def visit_Assert(self, node):
        return None  # Some processing here

    def visit_FunctionDef(self, node):
        i = len(node.body) - 1
        while i >= 0:
            line_node = node.body[i]
            if isinstance(line_node, ast.Assert):
                line_node = self.visit_Assert(line_node)
                if line_node is None:
                    del node.body[i]
                    i -= 1
                    continue
            else:
                for test_node in ast.walk(line_node):
                    # More processing here
                    pass
            i -= 1
        return node

Visiting a assert statement within a functiondef using Python ast

1 Answers1