CODEWITHSUNDEEP
pythonfunctionclassnestedabstract-syntax-tree
JJ_29
JJ_29

Reputation: 53

Extracting all classes and functions from an unknown python file

I am using AST library to extract all the details of a file like this,

import ast
file = open("TestFile.py", "r")
f = file.read()
p = ast.parse(f)
classFunc = [node.name for node in ast.walk(p) if isinstance(node, ast.ClassDef) or isinstance(node, ast.FunctionDef)]
print classFunc

this gives me an output,

['adf', 'A', 'message', 'dghe', '__init__', 'mess', 'B', 'teqwtdg']

Here, 'adf' and 'A' are the main classes, 'message' and 'dghe' are the functions under 'adf', 'init' and 'mess' are the functions under A, 'B' is a class under 'A' and 'teqwtdg' is a function under 'B'.

So, now my task is to write a python file where I create the class instances and call these functions(these are from an unknown file). I want to arrange this list such that I can easily know which are the main classes, which are the sub classes and which function comes under which class. How can I do this?

Upvotes: 1

Views: 1249

Answers (1)

gsb22
gsb22

Reputation: 2180

Following piece of code would traverse the file and create an object in the hierarchy. 

import ast
import pprint


def create_py_object(node_to_traverse, current_object):
    for node in node_to_traverse.body:
        if isinstance(node, ast.ClassDef):
            current_object.append({node.name: []})
            create_py_object(node, current_object[-1:][0][node.name])
        if isinstance(node, ast.FunctionDef):
            current_object.append({node.name: 'func'})


file = open("TestFile.py", "r")
f = file.read()
node_to_traverse = ast.parse(f)
py_file_structure = []

create_py_object(node_to_traverse, py_file_structure)
pprint.pprint(py_file_structure)

Output : 

[{'adf': [{'message': 'func'}, {'dghe': 'func'}]},
 {'A': [{'__init__': 'func'}, {'mess': 'func'}, {'B': [{'teqwtdg': 'func'}]}]}]

Upvotes: 1

Related Questions