Home > Mobile >  Adding visit function to AST visitor prevents other methods from firing
Adding visit function to AST visitor prevents other methods from firing

Time:05-04

Using the ast module, I run into a weird issue:

Code

import ast
from _ast import Try

code = """
def my_add(a, b):
    try:
        return a   b
    except TypeError:
        return a   b
"""

class Visitor(ast.NodeVisitor):
    def visit_Try(self, node: Try):
        print(node)
        return node


Visitor().visit(ast.parse(code))

Output

<ast.Try object at 0xXXX>

Here, Visitor().visit() (a ast.NodeVisitor subclass) prints <ast.Try object at 0xXXX>, but if I add a visit_FunctionDef method to the Visitor class:

Code with visit_FunctionDef

import ast
from _ast import Try

code = """
def my_add(a, b):
    try:
        return a   b
    except TypeError:
        return a   b
"""

class Visitor(ast.NodeVisitor):
    def visit_Try(self, node: Try):
        print(node)
        return node

    def visit_FunctionDef(self, node: ast.FunctionDef):
        print(node.name)
        return node

Visitor().visit(ast.parse(code))

Output

my_add

Expected Output

my_add
<ast.Try object at 0xXXX>

Now the only printed line is my_add, where I'd expect two lines, my_add, then <ast.Try object at 0xXXX>.

I checked and the visit_Try method seems okay, as it seems to respect what I can read here: class ast.NodeVisitor.visit. It isn't deprecated either.

I don't understand this behaviour. How can I get the desired output?

Same behaviour in Python 3.9 and 3.10.

CodePudding user response:

You need to add a call to .generic_visit() so that the children of the function node are visited. The children of a node aren't visited by default.

So, your visit_FunctionDef should look like:

def visit_FunctionDef(self, node: ast.FunctionDef):
    print(node.name)
    self.generic_visit(node)
    return node

This outputs:

my_add
<_ast.Try object at 0x7fbce58dd250>
  • Related