Extract classes and functions from Python files

pyextract.py

import argparse
import json
import ast
from ast import parse,walk
from astor import to_source

def extract_functions(node,is_class=False):
    """Recursively extract function names and code from an AST node."""
    functions = {}
    for child in walk(node):
        if isinstance(child, ast.FunctionDef):
            if is_class:
                functions[child.name] = to_source(child).strip()
            else:
                functions[child.name] = to_source(child).strip()
    return functions

def extract_classes(node):
    """Recursively extract class names and functions from an AST node."""
    classes = {}
    tmp_class_functions = {}
    independent_functions = {}
    for child in walk(node):
        if isinstance(child, ast.ClassDef):
            print("Is class")
            classes[child.name] = {"functions": extract_functions(child,is_class=True)}
            # Merge tmp_class_functions with new extracted functions
            tmp_class_functions = {**tmp_class_functions, **classes[child.name]["functions"]}
        # Check if the node is a function definition and not under a class
        elif isinstance(child, ast.FunctionDef) and child.name not in tmp_class_functions:
            independent_functions[child.name] = to_source(child).strip()
    return classes, independent_functions

def extract_info(file_name):
    """Extract class and function information from a Python file."""
    with open(file_name, "r") as f:
        source = f.read()
    root = parse(source)
    classes, independent_functions = extract_classes(root)
    return {"classes": classes, "independent_functions": independent_functions}

def main():
    parser = argparse.ArgumentParser(description='Extract class and function information from a Python file.')
    parser.add_argument('--in', dest='input_file', type=str, required=True, help='path to the input Python file')
    parser.add_argument('--out', dest='output_file', type=str, required=True, help='path to the output JSON file')
    args = parser.parse_args()

    data = extract_info(args.input_file)
    json_data = json.dumps(data, indent=4)

    with open(args.output_file, "w") as f:
        f.write(json_data)
    print(f'Data written to {args.output_file}')

if __name__ == "__main__":
    main()

Plan here is to link this to OpenAI's insert/completion endpoint to automatically add docstrings.

HTTP POST request with code file → Parse into JSON via this code → Find beginning of function/class → Insert Docstring quotes and [insert] tag → Send to davinci codec insert endpoint → Insert docstring → Return commented text.