2

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.

BrokenBenchmark
  • 18,126
  • 7
  • 21
  • 33
Dorian Turba
  • 3,260
  • 3
  • 23
  • 67

1 Answers1

1

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>
BrokenBenchmark
  • 18,126
  • 7
  • 21
  • 33
  • "Note that child nodes of nodes that have a custom visitor method won’t be visited unless the visitor calls generic_visit() or visits them itself." I didn't read the `generic_visit` documentation as I "guess" I didn't need it. Next time, I will try to open my eyes :D – Dorian Turba May 03 '22 at 17:10
  • 1
    Don't be so harsh on yourself, happens to all of us. :) Have a good one! – BrokenBenchmark May 03 '22 at 17:12