Automatically generate elaborate high quality docstrings for your Python file

auto_docstring.py

#!/usr/bin/env python
# coding: utf-8

import argparse
import os
import time

import openai
from tqdm import tqdm


def generate_docstring(file, openai_token):
    openai.api_key = openai_token

    with open(file) as f:
        lines = f.read()

    functions = lines.split('def ')[1:]
    functions[-1] = functions[-1].split('if __name__')[0]

    i = 1

    for func in tqdm(functions):

        func_tokens = 'def ' + func.replace('        ', '    ').strip()
        if '"""' in func_tokens:
            func_tokens = ''.join(
                func_tokens.split('"""')[0::2]).strip().replace('\n\n', '\n')

        func_name = func_tokens.split('(')[0].split('def ')[1]
        print(f'# >>>>>>>>>>>>>>>>>>>> METHOD/FUNCTION: {func_name}\n')

        prompt_tokens = len(func_tokens) / 2
        max_tokens = int(4000 - prompt_tokens)

        response = openai.Completion.create(
            engine="code-davinci-002",
            prompt=f"# Python 3.7\n \n{func_tokens}\n    \n# An elaborate, "
            "high quality docstring for the above function in google "
            "format:\n\"\"\"",
            temperature=0,
            max_tokens=max_tokens,
            top_p=1,
            frequency_penalty=0,
            presence_penalty=0.2,
            stop=["#", "\"\"\""])

        try:
            resp = response['choices'][0]['text'].strip()
        except KeyError as e:
            print(f'ERROR!: {e}')

        if resp.startswith('    '):
            resp = resp[5:]
        print('    """' + resp.strip())
        print('    """')
        print(f'\n{"-" * 80}\n')

        if i != len(functions):
            time.sleep(20)
            i += 1


def opts() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('-f',
                        '--file',
                        help='Path to the input file',
                        type=str,
                        required=True)
    parser.add_argument('-t',
                        '--token',
                        help='OpenAI token',
                        type=str,
                        default=os.getenv('OPENAI_API_KEY'))
    return parser.parse_args()


if __name__ == '__main__':
    args = opts()
    generate_docstring(file=args.file, openai_token=args.token)

usage.md

1. Sign up or login to OpenAI.
2. Go to your API keys page.
3. Create a new secret key and copy it.
4. In a terminal shell, run:

export OPENAI_TOKEN='xxxxxxxxxxx'

5. Install the required packages.

pip install openai tqdm

6. Now you can use the auto_docstring.py script to automatically generate elaborate docstrings for your Python files by running:

python auto_docstring.py -f demo.py -t "$OPENAI_TOKEN"