Build a Code Documentation Generator with AI

Build a Code Documentation Generator with AI

You’ve spent hours writing a complex function, but the moment you open it up three months later, you’re staring at a wall of logic with no idea what it does. That’s the silent killer of developer productivity: outdated or missing documentation. While writing docs by hand is tedious and often skipped, AI can now generate accurate, human-readable documentation in seconds—turning your codebase into a self-explaining library. Today, you’ll build your own AI-powered Code Documentation Generator in Python that you can run locally and integrate into your workflow immediately.

Why AI Documentation Actually Works

Traditional documentation tools rely on static patterns or regex, which fail when code gets complex. AI, specifically Large Language Models (LLMs), understands context, intent, and structure. It doesn’t just see def calculate(x, y); it sees why those variables exist and how they interact.

According to industry analysis, tools like Stenography and GitHub Copilot are already the practical starting points for populating missing docstrings, while Mintlify excels at publishing external docs [1]. But instead of relying on a paid service, building your own generator gives you full control over the output format, the model used, and the integration logic.

The Architecture: From Code to Docs

Before writing code, let’s map out the pipeline. A robust generator follows four key steps:

1. Code Parsing: Extract structured metadata (function signatures, type hints, class names) from raw source files.
2. Context Enrichment: Feed that metadata into an LLM along with the code body.
3. Generation: The LLM produces docstrings, parameter descriptions, and usage examples.
4. Injection: Merge the generated docs back into the source file or output a Markdown report.

This approach mirrors how advanced tools like CodeGPT operate: you highlight code, the system analyzes it, and the AI instantly generates comprehensive documentation with return values and best practices [2].

Building Your Python Generator

Let’s build a working generator using Python and the openai library (or any compatible LLM provider). This script will:

• Parse a Python file.
• Identify functions and classes.
• Generate docstrings using an LLM.
• Save the updated code.

Prerequisites

You’ll need:

• Python 3.8+
• An API key from an LLM provider (e.g., OpenAI, Anthropic, or a local model via Ollama).
• The openai and ast libraries.

Install dependencies:

pip install openai

The Code

Here’s a complete, runnable script. Replace YOUR_API_KEY and adjust the model name if you’re using a different provider.

import ast
import openai
import sys

# Configure your LLM client
openai.api_key = "YOUR_API_KEY"
MODEL = "gpt-4o-mini"  # Or use "ollama/local-model" if using Ollama

def extract_functions(code: str) -> list:
    """Parse Python code and extract function/class definitions."""
    tree = ast.parse(code)
    functions = []
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.ClassDef)):
            functions.append({
                "name": node.name,
                "args": [arg.arg for arg in node.args.args],
                "body": ast.get_source_segment(code, node)
            })
    return functions

def generate_docstring(func_info: dict) -> str:
    """Ask the LLM to generate a docstring for a function."""
    prompt = f"""
    Generate a professional Google-style docstring for this Python function:

    Name: {func_info['name']}
    Arguments: {func_info['args']}
    Code:
    {func_info['body']}

    Include:
    - A clear one-sentence description
    - Args section with types and descriptions
    - Returns section (if applicable)
    - One usage example
    """

    response = openai.ChatCompletion.create(
        model=MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.3
    )

    return response.choices[0].message.content.strip()

def inject_docstrings(code: str, docstrings: dict) -> str:
    """Insert generated docstrings into the original code."""
    tree = ast.parse(code)
    lines = code.splitlines()

    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef) and node.name in docstrings:
            # Find the line number of the function definition
            line_no = node.lineno - 1
            # Insert docstring after the function definition line
            docstring = f'    """{docstrings[node.name]}"""\n'
            lines.insert(line_no + 1, docstring)

    return "\n".join(lines)

def main(filepath: str):
    with open(filepath, "r") as f:
        code = f.read()

    functions = extract_functions(code)
    print(f"Found {len(functions)} functions/classes to document.")

    docstrings = {}
    for func in functions:
        print(f"Generating docs for: {func['name']}")
        docstrings[func['name']] = generate_docstring(func)

    updated_code = inject_docstrings(code, docstrings)

    # Save to a new file
    output_path = filepath.replace(".py", "_documented.py")
    with open(output_path, "w") as f:
        f.write(updated_code)

    print(f"✅ Documentation saved to {output_path}")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python doc_gen.py your_script.py")
        sys.exit(1)
    main(sys.argv[1])

How to Run It

1. Save this as doc_gen.py.
2. Create a test file, e.g., calc.py:

   def add(a, b):
       return a + b

1. Run: python doc_gen.py calc.py
2. Check calc_documented.py for the auto-generated docstring.

This script mirrors the workflow of tools like CodeDocAI, which let you paste code, adjust parameters, and generate Markdown docs instantly [5].

Making It Production-Ready

The script above is a great starting point, but real-world codebases need more. Here’s how to level it up:

1. Batch Processing for Entire Repositories

Instead of processing one file, traverse your entire Git repository. Tools like the one in the video tutorial scan every function and class, building a complete inventory before feeding it to the LLM in batch [3]. This is faster and ensures context across modules.

2. Customization and Style Matching

Not all teams use Google-style docs. Some prefer NumPy, Sphinx, or custom formats. You can add a --style flag to your script and pass the desired format to the LLM prompt. Mintlify allows users to edit and refine auto-generated docs, adding examples or adjusting formatting [11].

3. CI/CD Integration

The most powerful use of AI docs is automation. Integrate your generator into your CI/CD pipeline so docs update automatically when code changes. Set triggers for API updates or new module additions [15]. This ensures your documentation stays synchronized with your codebase without manual effort [3].

4. Human Review Loop

AI isn’t perfect. Always include a review step. GitHub Copilot and AskCodi both recommend reviewing generated docstrings to ensure clarity and accuracy before committing [11][1]. You can even add a --review flag that outputs the docs to a separate file for team validation.

Choosing the Right Model

Your generator’s quality depends heavily on the LLM you choose:

• GPT-4o: Highest accuracy, best for complex logic.
• Claude 3: Excellent for long-context code understanding.
• Local Models (Ollama/Llama): Free, private, but may need tuning.

For production, start with a cloud model for reliability, then migrate to local models if cost or privacy becomes a concern.

What’s Next?

You now have a working AI Code Documentation Generator that you can run today. But the real magic happens when you scale it.

• Extend to other languages: Add parsers for TypeScript, Go, or Java.
• Build a web UI: Use the CodeDocAI web interface pattern where users paste code or upload files [5].
• Generate full docs: Beyond docstrings, create READMEs, architecture overviews, and UML diagrams like DocuWriter.ai does [4].
• Publish automatically: Use MKDocs or Sphinx to build a polished API site from your documented source [3].

Documentation shouldn’t be a burden. With AI, it becomes a seamless part of your development flow. **Try running

---

If you found this helpful, consider buying me a coffee ☕ — it keeps these articles coming!

Also check out my AI tools collection: AI 次元世界 — free AI tools for developers.

---

💡 Related: **Content Creator Ultimate Bundle (Save 33%)* — $29.99*