Genius

docugenius.genius.genius

Genius(docstring_format='google', skip_raises=False, skip_returns=False, skip_examples=False, **kwargs)

Bases: ABC

An abstract class for generating docstrings for Python code.

Parameters:

Name	Type	Description	Default
docstring_format	Literal['google', 'numpy', 'sprinx']	The format of the docstring. Defaults to "google".	'google'
skip_raises	bool	Whether to include exceptions in the docstring. Defaults to False.	False
skip_returns	bool	Whether to include return values in the docstring. Defaults to False.	False
skip_examples	bool	Whether to include examples in the docstring. Defaults to False.	False

Raises:

Type	Description
ValueError	If the provided docstring_format is not valid.

Source code in docugenius/genius/genius.py

def __init__(
    self,
    docstring_format: Literal["google", "numpy", "sprinx"] = "google",
    skip_raises: bool = False,
    skip_returns: bool = False,
    skip_examples: bool = False,
    **kwargs,
):
    self.docstring_format = docstring_format
    if docstring_format not in ["google", "numpy", "sprinx"]:
        raise ValueError(
            f"Invalid docstring format {docstring_format}. Must be one of 'google', 'numpy', or 'sprinx'."
        )
    self.skip_raises = skip_raises
    self.skip_returns = skip_returns
    self.skip_examples = skip_examples
    super().__init__()

clean_llm_output(output)

Clean the output from the LLM model by removing any leading or trailing whitespace.

Parameters:

Name	Type	Description	Default
output	str	The output from the LLM model.	required

Returns:

Name	Type	Description
str	str	The cleaned output.

Source code in docugenius/genius/genius.py

def clean_llm_output(self, output: str) -> str:
    """
    Clean the output from the LLM model by removing any leading or trailing whitespace.

    Parameters:
        output (str): The output from the LLM model.

    Returns:
        str: The cleaned output.
    """
    import re

    output = re.findall(r"```generated-python-code\n(.*?)\n```", output, re.DOTALL)
    if output:
        output = output[0]
    else:
        output = ""

    return output

generate(code)

Generate a docstring for the given code.

Parameters:

Name	Type	Description	Default
code	str	The code to generate a docstring for.	required

Returns:

Name	Type	Description
str	str	The code with the generated docstring added.

Source code in docugenius/genius/genius.py

def generate(self, code: str) -> str:
    """
    Generate a docstring for the given code.

    Parameters:
        code (str): The code to generate a docstring for.

    Returns:
        str: The code with the generated docstring added.
    """
    if not code or is_skippable(code):
        return code

    output = self._generate_docstring(code)
    return self.clean_llm_output(output)

docugenius.genius.openai_genius

OpenAIGenius(model, docstring_format='google', skip_raises=False, skip_returns=False, skip_examples=False)

Bases: Genius

A class to generate docstrings using OpenAI's API.

Parameters:

Name	Type	Description	Default
model	str	The model to use for generating docstrings.	required
docstring_format	Literal['google', 'numpy', 'sprinx']	The format of the docstring. Defaults to "google".	'google'
skip_raises	bool	Whether to include exceptions in the docstring. Defaults to False.	False
skip_returns	bool	Whether to include return values in the docstring. Defaults to False.	False
skip_examples	bool	Whether to include examples in the docstring. Defaults to False.	False

Raises:

Type	Description
ValueError	If the provided model is not valid.

Examples:

>>> genius = OpenAIGenius(model="gpt-4o-mini")
>>> print(genius.docstring_format)
google

Initialize the OpenAIGenius instance.

Parameters:

Name	Type	Description	Default
model	str	The model to use for generating docstrings.	required
docstring_format	Literal['google', 'numpy', 'sprinx']	The format of the docstring. Defaults to "google".	'google'
skip_raises	bool	Whether to include exceptions in the docstring. Defaults to False.	False
skip_returns	bool	Whether to include return values in the docstring. Defaults to False.	False
skip_examples	bool	Whether to include examples in the docstring. Defaults to False.	False

Raises:

Type	Description
ValueError	If the provided model is not valid.

Examples:

>>> genius = OpenAIGenius(model="gpt-4o-mini")

Source code in docugenius/genius/openai_genius.py

def __init__(
    self,
    model: str,
    docstring_format: Literal["google", "numpy", "sprinx"] = "google",
    skip_raises: bool = False,
    skip_returns: bool = False,
    skip_examples: bool = False,
):
    """
    Initialize the OpenAIGenius instance.

    Args:
        model (str): The model to use for generating docstrings.
        docstring_format (Literal["google", "numpy", "sprinx"], optional): The format of the docstring. Defaults to "google".
        skip_raises (bool, optional): Whether to include exceptions in the docstring. Defaults to False.
        skip_returns (bool, optional): Whether to include return values in the docstring. Defaults to False.
        skip_examples (bool, optional): Whether to include examples in the docstring. Defaults to False.

    Raises:
        ValueError: If the provided model is not valid.

    Examples:
        >>> genius = OpenAIGenius(model="gpt-4o-mini")
    """
    import openai

    super().__init__(docstring_format, skip_raises, skip_returns, skip_examples)
    self._client = openai.OpenAI()
    self.model = model

clean_llm_output(output)

Clean the output from the LLM model by removing any leading or trailing whitespace.

Parameters:

Name	Type	Description	Default
output	str	The output from the LLM model.	required

Returns:

Name	Type	Description
str	str	The cleaned output.

Source code in docugenius/genius/genius.py

def clean_llm_output(self, output: str) -> str:
    """
    Clean the output from the LLM model by removing any leading or trailing whitespace.

    Parameters:
        output (str): The output from the LLM model.

    Returns:
        str: The cleaned output.
    """
    import re

    output = re.findall(r"```generated-python-code\n(.*?)\n```", output, re.DOTALL)
    if output:
        output = output[0]
    else:
        output = ""

    return output

generate(code)

Generate a docstring for the given code.

Parameters:

Name	Type	Description	Default
code	str	The code to generate a docstring for.	required

Returns:

Name	Type	Description
str	str	The code with the generated docstring added.

Source code in docugenius/genius/genius.py

def generate(self, code: str) -> str:
    """
    Generate a docstring for the given code.

    Parameters:
        code (str): The code to generate a docstring for.

    Returns:
        str: The code with the generated docstring added.
    """
    if not code or is_skippable(code):
        return code

    output = self._generate_docstring(code)
    return self.clean_llm_output(output)