Search code examples
pythonpython-3.xabstract-syntax-tree

Python ast: decide whether a FunctionDef is inside a ClassDef or not


I'd like to build an ast from a Python source code, then get specific information from the ast. I faced the following problem: while walking through a ClassDef's body is viable, how I decide if a method is inside the class or not.

The code I build ast from:

class A:
    def foo(self):
        pass


def foo(self):
    pass

In this example I will hit all the foos but I will not able to tell if it is from the class or not (as they have the same set of parameters, named badly, but the code can be interpreted).

    def build_ast(self):
        with open(self.path, 'r', encoding='utf-8') as fp:
            tree = ast.parse(fp.read())
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    print(ast.dump(node))
                    # access the parent if it has

Solution

  • I'm not totally satisfied with my final solution, but apparently it works for Python 3.8.3:

    As I experienced, ast.walk traverses ClassDef nodes before FunctionsDef nodes.

    def build_ast(self):
        with open(self.path, 'r', encoding='utf-8') as fp:
            tree = ast.parse(fp.read())
    
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    if hasattr(node, "parent"):
                        print(node.parent.name, node.name)
                    else:
                        print(node.name, "is not in a class.")
                if isinstance(node, ast.ClassDef):
                    for child in node.body:
                        if isinstance(child, ast.FunctionDef):
                            child.parent = node