Uses GPT to generate docstrings for functions in python files automagically

docstringer.py

#!/usr/bin/env python3
# By Arturo 'Buanzo' Busleiman
# YMMV
import openai
import ast
import astunparse
import sys
import shutil
from itertools import takewhile

def run_gpt(role, prompt, model='gpt-3.5-turbo'):
    response = openai.ChatCompletion.create(model=model,
                                            messages=[{"role": "system", "content": role},
                                                      {"role": "user", "content": prompt}
                                                      ])
    respuesta = response["choices"][0]["message"]["content"]
    return respuesta

def insert_docstrings(file_path):
    with open(file_path, 'r') as f:
        code = f.read()

    tree = ast.parse(code)

    doc_insert_positions = []

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef):
            func_code = astunparse.unparse(node)
            role = """you generate docstrings that describe the purpose, arguments and returns for python functions.
you are given an entire function implementation, and you return ONLY the actual docstring, without the triple quotes.
limit each line to less than 80 characters."""
            docstring_content = run_gpt(role, func_code)
            start_line = node.lineno
            doc_insert_positions.append((start_line, docstring_content))

    doc_insert_positions.sort(key=lambda x: x[0], reverse=True)

    lines = code.split('\n')

    for line_number, docstring_content in doc_insert_positions:
        leading_whitespace = ''.join(takewhile(str.isspace, lines[line_number - 1]))
        indented_docstring = leading_whitespace + '    """\n'
        indented_docstring += leading_whitespace + '    ' + docstring_content.replace("\n", "\n" + leading_whitespace + '    ') + '\n'
        indented_docstring += leading_whitespace + '    """'
        lines.insert(line_number, indented_docstring)

    new_code = '\n'.join(lines)
    return new_code

def main():
    if len(sys.argv) < 2:
        print("Usage: python docstringer.py <python_file.py>")
        return

    file_path = sys.argv[1]
    shutil.copy(file_path, f"{file_path}.bak")  # Backup the original file

    new_code = insert_docstrings(file_path)

    with open(file_path, 'w') as f:
        f.write(new_code)

if __name__ == "__main__":
    main()