Codegen’s GraphSitter library provides powerful APIs for editing docstrings and adding documentation to your codebase. This guide will walk you through the process of modifying existing docstrings and adding new documentation to functions, classes, and modules.

Common use cases include:

  • Updating outdated docstrings
  • Adding missing parameter and return type descriptions
  • Standardizing documentation format across the codebase
  • Adding examples to function docstrings

​
set_docstring

Editing docstrings and adding documentation involves accessing and modifying the docstring property of various code elements. Here’s a basic example:

By default, set_docstring tries to intelligently format the docstring to match the existing format. This can be turned on and off using auto_format and clean_format flags.

​
Determine Documentation Coverage

To see your current documentation coverage, you can iterate through all symbols of interest and count the number of docstrings:

python
# Initialize counters
total_functions = 0
functions_with_docs = 0
total_classes = 0
classes_with_docs = 0

# Check functions
for function in codebase.functions:
    total_functions += 1
    if function.docstring:
        functions_with_docs += 1

# Check classes
for cls in codebase.classes:
    total_classes += 1
    if cls.docstring:
        classes_with_docs += 1

# Calculate percentages
func_coverage = (functions_with_docs / total_functions * 100) if total_functions > 0 else 0
class_coverage = (classes_with_docs / total_classes * 100) if total_classes > 0 else 0

# Print results with emojis
print("\n📊 Documentation Coverage Report:")
print(f"\n📝 Functions:")
print(f"  • Total: {total_functions}")
print(f"  • Documented: {functions_with_docs}")
print(f"  • Coverage: {func_coverage:.1f}%")

print(f"\n📚 Classes:")
print(f"  • Total: {total_classes}")
print(f"  • Documented: {classes_with_docs}")
print(f"  • Coverage: {class_coverage:.1f}%")

print(f"\n🎯 Overall Coverage: {((functions_with_docs + classes_with_docs) / (total_functions + total_classes) * 100):.1f}%")

​
Leveraging AI for Generating Documentation

For non-trivial codebases, it can be challenging to achieve full documentation coverage.

The most efficient way to edit informative docstrings is to use codebase.ai to generate docstrings, then use the set_docstring method to update the docstring.

python
# Import datetime for timestamp
from datetime import datetime

# Get current timestamp
timestamp = datetime.now().strftime("%B %d, %Y")

print("📚 Generating and Updating Function Documentation")

# Process all functions in the codebase
for function in codebase.functions:
    current_docstring = function.docstring()

    if current_docstring:
        # Update existing docstring to be more descriptive
        new_docstring = codebase.ai(
            f"Update the docstring for {function.name} to be more descriptive and comprehensive.",
            target=function
        )
        new_docstring += f"\n\nUpdated on: {timestamp}"
    else:
        # Generate new docstring for function
        new_docstring = codebase.ai(
            f"Generate a comprehensive docstring for {function.name} including parameters, return type, and description.",
            target=function
        )
        new_docstring += f"\n\nCreated on: {timestamp}"

    # Set the new or updated docstring
    function.set_docstring(new_docstring)
More info on using AI with Codebase can be found here.

​
Adding Explicit Parameter Names and Types

Alternatively, you can also rely on deterministic string formatting to edit docstrings.

To add “Google-style” parameter names and types to a function docstring, you can use the following code snippet:

python
# Iterate through all functions in the codebase
for function in codebase.functions:
    # Skip if function already has a docstring
    if function.docstring:
        continue

    # Build parameter documentation
    param_docs = []
    for param in function.parameters:
        param_type = param.type.source if param.is_typed else "Any"
        param_docs.append(f"    {param.name} ({param_type}): Description of {param.name}")

    # Get return type if present
    return_type = function.return_type.source if function.return_type else "None"

    # Create Google-style docstring
    docstring = f'''"""
    Description of {function.name}.

    Args:
{chr(10).join(param_docs)}

    Returns:
        {return_type}: Description of return value
    """'''

    # Set the new docstring
    function.set_docstring(docstring)

Was this page helpful?