Reference for llmcode/utils/document.py

This module provides functionality for documenting Python scripts.

It includes functions to parse Python scripts, extract classes and functions, and generate or update docstrings using a completion model.

Functions: doc_element: Generates a docstring for a given element using a completion model. doc_python_file: Documents a Python file by generating or updating docstrings for its elements.

Author: Francisco Javier Gañán License File: https://github.com/javierganan99/LLMCode/blob/main/LICENSE

llmcode.utils.document.doc_element

doc_element(element, prompt, get_completion)

Processes a specific element by generating a completion based on a provided prompt.

This function utilizes a completion generator, replaces a specific key in the prompt with the given element, and executes the generation process within a specified timeout limit.

Parameters:

Name	Type	Description	Default
element	str	The specific element to be processed and substituted in the prompt.	required
prompt	str	The prompt template that includes a placeholder for the element.	required
get_completion	callable	A function responsible for generating the completion based on the modified prompt.	required

Returns:

Type	Description
any	The result of the completion generation, which can vary depending on the
	implementation of the get_completion function.

Source code in llmcode/utils/document.py

def doc_element(element, prompt, get_completion):
    """
    Processes a specific element by generating a completion based on a provided prompt.

    This function utilizes a completion generator, replaces a specific key in the prompt with
    the given element, and executes the generation process within a specified timeout limit.

    Args:
        element (str): The specific element to be processed and substituted in the prompt.
        prompt (str): The prompt template that includes a placeholder for the element.
        get_completion (callable): A function responsible for generating the completion based
            on the modified prompt.

    Returns:
        (any): The result of the completion generation, which can vary depending on the
        implementation of the get_completion function.
    """
    return run_with_timeout(
        get_completion,
        args=(
            read_content(prompt).replace(custom_params.query_completion_key, element),
        ),
        timeout=custom_params.completion_timeout,
    )

---

llmcode.utils.document.doc_python_file

doc_python_file(script, prompts, elements2doc=custom_params.elements2doc, overwrite=custom_params.overwrite, get_completion=get_completion, stop_flag=Event(), TODO_message='# TODO: Document this ELEMENT on your own. Could not be documented by the model.')

Documents Python files by generating docstrings for classes and functions.

This function parses a given Python script, extracts the defined classes and functions, and generates docstrings for them using specified prompts and a completion function. It also handles existing docstrings based on the overwrite parameter and can stop the documentation process if triggered.

Parameters:

Name	Type	Description	Default
script	str	The path to the Python script file to be documented.	required
prompts	dict	A dictionary with the prompts to be used for generating docstrings.	required
elements2doc	list	A list of elements (classes or functions) to document. Defaults to custom_params.elements2doc.	elements2doc
overwrite	bool	Flag indicating whether to overwrite existing docstrings. Defaults to custom_params.overwrite.	overwrite
get_completion	callable	A function to get the completion responses. Defaults to get_completion.	get_completion
stop_flag	Event	An event flag to signal stopping the documentation process. Defaults to Event().	Event()
TODO_message	str	Message template to indicate that an element could not be documented. Defaults to "# TODO: Document this ELEMENT on your own. Could not be documented by the model."	'# TODO: Document this ELEMENT on your own. Could not be documented by the model.'

Source code in llmcode/utils/document.py

def doc_python_file(
    script,
    prompts,
    elements2doc=custom_params.elements2doc,
    overwrite=custom_params.overwrite,
    get_completion=get_completion,
    stop_flag=Event(),
    TODO_message="# TODO: Document this ELEMENT on your own. Could not be documented by the model.",
):
    """
    Documents Python files by generating docstrings for classes and functions.

    This function parses a given Python script, extracts the defined classes and functions, and generates docstrings
    for them using specified prompts and a completion function. It also handles existing docstrings based on the
    overwrite parameter and can stop the documentation process if triggered.

    Args:
        script (str): The path to the Python script file to be documented.
        prompts (dict): A dictionary with the prompts to be used for generating docstrings.
        elements2doc (list, optional): A list of elements (classes or functions) to document. Defaults to
            custom_params.elements2doc.
        overwrite (bool, optional): Flag indicating whether to overwrite existing docstrings. Defaults to
            custom_params.overwrite.
        get_completion (callable, optional): A function to get the completion responses. Defaults to get_completion.
        stop_flag (Event, optional): An event flag to signal stopping the documentation process. Defaults to
            Event().
        TODO_message (str, optional): Message template to indicate that an element could not be documented.
            Defaults to "# TODO: Document this ELEMENT on your own. Could not be documented by the model."
    """
    script = str(script)
    try:
        parsed = parse_python(script)
        if parsed is None:
            LOGGER.info(
                "%s\r❌Check the script {script}. It is not properly coded and can not be documented.",
                ANSI_CODE["red"],
            )
            return None
        (
            classes_and_functions,
            script_content,
        ) = extract_functions_and_classes_from_python_tokens(parsed)
    except Exception as e:
        LOGGER.info(
            "%s\r❌Check the script %s. The following error occurred: %s",
            ANSI_CODE["red"],
            script,
            e,
        )

    def add_msg(msg, where, element, script_content):
        n_spaces = 0
        for ch in element[where - 1 :: -1]:
            if ch == " ":
                n_spaces += 1
            else:
                break
        new_element = element[:where] + msg + "\n" + " " * n_spaces + element[where:]
        return script_content.replace(element, new_element)

    for e_type, elements in classes_and_functions.items():
        for element, e_name, dosctring_idx in elements:
            if (
                elements2doc is not None and e_name not in elements2doc
            ):  # Filter the elements to document
                continue
            if stop_flag.is_set():
                LOGGER.info(
                    "%s\r Ended during the documentation of script %s. Interrupted by SIGINT.%s",
                    ANSI_CODE["yellow"],
                    script,
                    ANSI_CODE["reset"],
                )
                return
            if e_type == "function" or e_type == "class":
                if element[dosctring_idx : dosctring_idx + 3] == '"""':
                    previous_docstring = re.search(r'"""(.*?)"""', element, re.DOTALL)
                    if (
                        previous_docstring and not overwrite
                    ):  # If element had docstring and we do not want to change it
                        continue
                    LOGGER.info(
                        "%s\r\n🤖 Generating docstring for %s %s...\n\n",
                        ANSI_CODE["reset"],
                        e_type,
                        e_name,
                    )
                else:
                    previous_docstring = None
                result = doc_element(element, prompts[e_type], get_completion)
                if result is not None:  # Query successfull
                    new_docstring = re.search(r'"""(.*?)"""', result, re.DOTALL)
                    if new_docstring:
                        new_docstring = new_docstring.group(1)
                        LOGGER.info("%s\r%s\n\n\n", ANSI_CODE["reset"], new_docstring)
                    else:
                        LOGGER.info(
                            "%s\r⚠ No docstring generated for %s %s...",
                            ANSI_CODE["yellow"],
                            e_type,
                            e_name,
                        )
                        script_content = add_msg(
                            TODO_message.replace("ELEMENT", e_type),
                            dosctring_idx,
                            element,
                            script_content,
                        )
                        continue
                    if previous_docstring:  # If element had docstring
                        previous_docstring = previous_docstring.group(1)
                        element_new = element.replace(
                            previous_docstring, new_docstring, 1
                        )
                        script_content = script_content.replace(element, element_new, 1)
                    else:  # If element did not have docstring
                        script_content = add_msg(
                            f'"""{new_docstring}"""',
                            dosctring_idx,
                            element,
                            script_content,
                        )
                else:  # Error in the query
                    LOGGER.info(
                        "%s\r❌ Error in the query for %s %s! No response provided...",
                        ANSI_CODE["red"],
                        e_type,
                        e_name,
                    )
                    script_content = add_msg(
                        TODO_message.replace("ELEMENT", e_type),
                        dosctring_idx,
                        element,
                        script_content,
                    )
    with open(script, "w", encoding="utf-8") as python_file:
        python_file.write(script_content)