utils.py

import ast
import glob
import os


def parse_file_content(string: str):
    first_break = string.find("---")
    last_break = string.rfind("---")
    if first_break == -1 and last_break == -1 or first_break == last_break:
        return None, None
    nl_after = string.find("\n", last_break)
    description = string[nl_after:]
    return string[first_break + 4 : last_break], description.strip("\n")


def parse_action(string: str):
    assert string.startswith("action:")
    idx = string.find("action_input=")
    if idx == -1:
        return string[8:], None
    return string[8 : idx - 1], string[idx + 13 :].strip("'").strip('"')


def extract_imports(file_contents):
    module_ast = ast.parse(file_contents)
    imports = []
    functions = [n for n in module_ast.body if isinstance(n, ast.FunctionDef)]
    classes = [n for n in module_ast.body if isinstance(n, ast.ClassDef)]
    for node in ast.walk(module_ast):
        if isinstance(node, ast.Import):
            for alias in node.names:
                imports.append(alias.name)
        elif isinstance(node, ast.ImportFrom):
            module_name = node.module
            for alias in node.names:
                name = alias.name
                if module_name:
                    imports.append(f"{module_name}.{name}")
                else:
                    imports.append(name)
    return imports, functions, classes


def read_python_module_structure(path):
    file_types = ["*.py"]
    code = []
    for file_type in file_types:
        code += glob.glob(os.path.join(path, "**", file_type), recursive=True)

    structure_prompt = "Files:\n"
    structure_prompt += "(listing all files and their functions and classes)\n\n"

    def get_file_name(i):
        return "./{}.py".format(i.replace(".", "/"))

    content = {}
    internal_imports_map = {}
    for fn in code:
        if os.path.basename(fn) == "gpt.py":
            continue
        with open(fn, "r") as f:
            content[fn] = f.read()
            imports, functions, classes = extract_imports(content[fn])
            internal_imports = list(
                {".".join(i.split(".")[:-1]) for i in imports if i.startswith("app.")}
            )
            internal_imports_map[fn] = [get_file_name(i) for i in internal_imports]
            structure_prompt += f"{fn}\n"
            for function in functions:
                structure_prompt += "  {}()\n".format(function.name, function.args.args)

            for class_ in classes:
                structure_prompt += "  {}\n".format(class_.name)
                methods = [n for n in class_.body if isinstance(n, ast.FunctionDef)]
                for method in methods:
                    structure_prompt += "  {}.{}()\n".format(class_.name, method.name)
            structure_prompt += "\n"

    return structure_prompt, content, internal_imports_map