Code Documentation Generator

Field	Value
Path	2 — AI-Augmented Engineering
Position	4 of 12
Difficulty	🟢 Local only
Estimated time	3 hours
AWS cost	None

---

Agent Pickup Instructions

# 1. Enter project
cd projects/p2-04-code-doc-generator

# 2. Create virtual environment
python -m venv .venv && source .venv/bin/activate

# 3. Install dependencies
pip install -r requirements.txt

# 4. Copy env file and add your API key
cp .env.example .env

# 5. Generate docs for a file
python src/main.py src/main.py --output README.md

# 6. Preview the diff against an existing README
python src/main.py src/ --diff

# 7. Run tests
pytest tests/ -v

Done when:

• python src/main.py <file> generates a markdown README section for every public function
• Private functions (leading _) are excluded from the output
• --diff flag shows a unified diff between generated docs and existing file
• All 4 tests pass
• Generated markdown has no unclosed code fences

---

What this project is

A read-only documentation generator that uses Python's ast module to extract function signatures, type hints, return types, and existing docstrings, then sends that structured metadata to an LLM to generate a README section and usage examples for each public function. It does not modify source files — it only generates documentation. A diff mode compares the generated output against any existing README, making it easy to see what documentation is new, changed, or missing.

---

What the learner achieves

"I built an AST-based documentation generator that extracts structured function metadata without executing the code, sends it to an LLM with a documentation-writing prompt, and produces a diff-able README — demonstrating that LLM tooling and static analysis are complementary, not competing."

---

Folder structure

p2-04-code-doc-generator/
├── src/
│   ├── main.py        # CLI: <file_or_dir> [--output README.md] [--diff]
│   ├── llm.py         # Provider-agnostic LLM wrapper
│   ├── ast_parser.py  # extract_functions -> list[FunctionInfo]
│   ├── generator.py   # generate_function_doc, generate_module_readme
│   └── differ.py      # show_diff between existing and generated
├── tests/
│   ├── fixtures/
│   │   ├── sample_module.py   # Public and private functions with type hints
│   │   └── existing_readme.md # Existing README to diff against
│   ├── test_ast_parser.py
│   ├── test_generator.py
│   └── test_differ.py
├── .env.example
├── requirements.txt
└── README.md

---

.env.example

# LLM provider: anthropic | openai | ollama
LLM_PROVIDER=anthropic

# Model to use
LLM_MODEL=

# API key for Anthropic
ANTHROPIC_API_KEY=

# API key for OpenAI (leave blank if using anthropic or ollama)
OPENAI_API_KEY=

# Ollama base URL
OLLAMA_BASE_URL=http://localhost:11434

# Maximum words per function doc section
DOC_MAX_WORDS=120

# Whether to include usage examples in generated docs (true | false)
INCLUDE_EXAMPLES=true

---

requirements.txt

anthropic==0.30.0
openai==1.35.0
python-dotenv==1.0.1
rich==13.7.1
pytest==8.2.2
pytest-mock==3.14.0

---

src/ — what to implement

src/llm.py

get_completion(prompt: str, system: str = "") -> str

• Standard interface: reads LLM_PROVIDER and LLM_MODEL from env
• Routes to Anthropic, OpenAI, or Ollama
• Edge cases: raise RuntimeError if API key is missing

---

src/ast_parser.py

FunctionInfo (dataclass)

@dataclass
class FunctionInfo:
    name: str
    params: list[dict]      # [{"name": str, "type": str | None, "default": str | None}]
    return_type: str | None
    docstring: str | None
    lineno: int
    is_async: bool
    decorators: list[str]

extract_functions(filepath: str) -> list[FunctionInfo]

• Inputs: path to a .py file
• Output: list of FunctionInfo for all public (non-underscore-prefixed) top-level and class-level functions
• Behavior:
  • Parse with ast.parse(source)
  • Walk FunctionDef and AsyncFunctionDef nodes
  • Skip if node.name.startswith("_")
  • For each param: extract arg.arg (name), ast.unparse(arg.annotation) if annotation exists (else None), and default value via ast.unparse if present
  • Return type: ast.unparse(node.returns) if node.returns else None
  • Docstring: ast.get_docstring(node)
  • Decorators: [ast.unparse(d) for d in node.decorator_list]
• Edge cases: syntax error → raise SyntaxError with filename in message; empty file → return []

extract_module_docstring(filepath: str) -> str | None

• Inputs: path to .py file
• Output: module-level docstring or None
• Behavior: parse with ast.parse; check if first statement is Expr(Constant(...))

---

src/generator.py

generate_function_doc(func: FunctionInfo) -> str

• Inputs: a FunctionInfo dataclass
• Output: markdown string for this function
• Behavior:
  • Build prompt: include function name, parameters with types and defaults, return type, existing docstring if present
  • System prompt: "You are a technical writer. Generate a concise markdown section for this Python function. Include: what it does, parameters table (Name | Type | Description), return value, and one usage example. Do not repeat the raw signature verbatim. Max {DOC_MAX_WORDS} words."
  • Call get_completion
  • Prepend ### {func.name} heading to result
• Edge cases: if existing docstring is present and > 50 words, use it directly without calling the LLM (mark as [from existing docstring])

generate_module_readme(filepath: str, functions: list[FunctionInfo]) -> str

• Inputs: filepath, list of FunctionInfo
• Output: full markdown README string
• Behavior:
  • Start with # <module_name> heading
  • Add module docstring section if present
  • For each function in functions, call generate_function_doc and concatenate
  • Add ## Quick reference table at top: Function | Parameters | Returns (generated from metadata, no LLM)
• Edge cases: if functions is empty, return a minimal README with module name and "No public functions found."

---

src/differ.py

show_diff(existing_path: str, generated: str) -> str

• Inputs: path to existing file (may not exist), generated content string
• Output: unified diff string
• Behavior: if existing file exists, read it; compute unified diff using difflib.unified_diff; return as string; if file does not exist, return diff against empty string (all additions)

print_diff(diff: str) -> None

• Inputs: diff string
• Output: prints with color — additions in green, removals in red (use rich for coloring)
• Edge cases: if diff is empty (no changes), print "No changes — generated docs match existing file."

---

src/main.py

CLI entry point using argparse.

Arguments:

• Positional: target — path to a Python file or directory
• --output <path> — write generated README to this path; if omitted, print to stdout
• --diff — show diff between generated output and existing --output file (requires --output)

Behavior:

1. Load .env
2. If target is a directory, find all .py files recursively
3. For each file: call extract_functions; log Parsed N public functions from <file>
4. Call generate_module_readme for each file
5. Concatenate all READMEs if multiple files
6. If --diff: call show_diff and print_diff; exit without writing
7. If --output: write to file; print Written to <path>
8. Otherwise: print to stdout

---

tests/ — what to test

Test 1 — AST parser extracts correct function names and params (test_ast_parser.py)

• Use tests/fixtures/sample_module.py which has a function def connect(host: str, port: int = 5432) -> None:
• Call extract_functions
• Assert connect is in the returned names
• Assert the host param has type == "str" and no default
• Assert the port param has type == "int" and default == "5432"

Test 2 — Private functions are skipped (test_ast_parser.py)

• sample_module.py has def _internal_helper():
• Call extract_functions
• Assert _internal_helper is NOT in returned names

Test 3 — Type hints captured correctly (test_ast_parser.py)

• sample_module.py has def process(data: list[str]) -> dict[str, int]:
• Call extract_functions
• Assert the return type string is "dict[str, int]"
• Assert the data param type string is "list[str]"

Test 4 — Generated markdown has no unclosed code fences (test_generator.py)

• Mock get_completion to return a response with a code block
• Call generate_function_doc with a minimal FunctionInfo
• Assert the output: count occurrences of ```; must be even (every opened fence is closed)
• Assert the output starts with ### (the heading is prepended)

Test 5 — Module readme has quick reference table (test_generator.py)

• Call generate_module_readme with 3 mock FunctionInfo objects (mock the LLM)
• Assert the result contains ## Quick reference
• Assert all 3 function names appear in the quick reference table

Test 6 — Diff shows additions when no existing file (test_differ.py)

• Call show_diff with a path that does not exist
• Assert every non-header line in the diff starts with +

---

README.md content

# Code Documentation Generator

Reads Python files with AST, extracts function signatures and type hints, and generates a README with LLM-written documentation for every public function.

## Quick start

```bash
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
cp .env.example .env   # add your API key

# Generate docs for a file
python src/main.py src/main.py

# Generate and write to README
python src/main.py src/ --output README.md

# Preview what would change
python src/main.py src/ --output README.md --diff

What it generates

For each public function:

• What it does (1-2 sentences)
• Parameters table (Name | Type | Description)
• Return value description
• One usage example

Rules

• Private functions (_name) are always skipped
• Functions with existing docstrings > 50 words skip the LLM call
• No source files are modified — read-only

Running tests

pytest tests/ -v

---

## GUIDE.md content

```markdown
# Build Guide — Code Documentation Generator

## Step 1 — AST extraction without executing code

The `ast` module lets you analyze Python without running it:

```python
import ast

with open(filepath) as f:
    source = f.read()

tree = ast.parse(source)

for node in ast.walk(tree):
    if isinstance(node, ast.FunctionDef):
        name = node.name
        return_type = ast.unparse(node.returns) if node.returns else None
        docstring = ast.get_docstring(node)
        params = [
            {
                "name": arg.arg,
                "type": ast.unparse(arg.annotation) if arg.annotation else None,
            }
            for arg in node.args.args
        ]

This is safer than exec or importlib — no side effects, no dependency installation required.

Step 2 — The documentation prompt

System prompt:

You are a technical writer. Document this Python function in markdown.
Include: one-sentence description, parameters table (Name | Type | Description),
return value, and one short usage example.
Do not repeat the function signature verbatim.
Max {DOC_MAX_WORDS} words.

User content:

Function: {name}
Parameters: {params_formatted}
Return type: {return_type}
Existing docstring: {docstring or "none"}

Step 3 — Quick reference table (no LLM)

Generate the quick reference table from metadata — no LLM call needed:

## Quick reference

| Function | Parameters | Returns |
|---|---|---|
| `connect` | host: str, port: int = 5432 | None |
| `query` | sql: str | list[dict] |

This is the first thing a reader sees. It's derived from AST, so it's always accurate.

Step 4 — Diff display

Use difflib.unified_diff:

import difflib

existing_lines = existing_content.splitlines(keepends=True)
generated_lines = generated_content.splitlines(keepends=True)
diff = difflib.unified_diff(existing_lines, generated_lines, fromfile="existing", tofile="generated")

Color with rich: iterate diff lines, print additions in green, removals in red.

Debugging tips

• If the LLM generates a code block but doesn't close it, count ``` occurrences — odd count = unclosed fence
• If type hints show as None when you expect a type, check that the file uses from __future__ import annotations; in that case ast.unparse may return string literals
• Test on your own src/main.py — dogfooding catches formatting issues quickly

How to talk about this in an interview

"Why AST instead of just reading the raw text?"

AST gives structured data: I can extract param names, types, defaults, and return types as typed fields, not strings I have to parse. The LLM gets structured input and produces structured output — no prompt hacking needed.

"What's the LLM doing that you couldn't do without it?"

The prose description: "This function establishes a connection to a PostgreSQL database and returns a cursor." That sentence requires understanding what host, port, and the return type mean together. Regex or templates can't do that.

"How do you handle existing docstrings?"

If there's a docstring longer than 50 words, I use it directly without an LLM call. This respects the developer's intent and saves cost. The LLM only fills gaps.

---

## Failure modes to handle

| Failure | Where | How to handle |
|---|---|---|
| File has syntax error | `ast_parser.py` | Raise `SyntaxError` with filename; `main.py` catches it, logs error, skips file |
| Function has no type hints | `ast_parser.py` | Set `type == None`; `generator.py` prompt says "type unknown" |
| LLM returns unclosed code fence | `generator.py` | Post-process: count ` ``` ` occurrences; if odd, append closing ` ``` ` |
| Target directory has no `.py` files | `main.py` | Print "No Python files found in <dir>." and exit 0 |
| `--diff` used without `--output` | `main.py` | Print "Error: --diff requires --output <path>" and exit 1 |
| Existing README is very large | `differ.py` | Diff is still computed; no truncation needed (difflib handles any size) |

---

## The metric this project measures

**What is measured:** Number of public functions documented, LLM calls made vs. skipped (docstring already present), and whether all code fences in output are balanced.

**Format (stdout):**

Parsed 8 public functions from src/main.py LLM calls: 6 | Skipped (existing docstring): 2 Written to README.md

**Target:** Every public function must appear in the output. Skipped-due-to-docstring count must be accurate (verify by checking source). All code fences in generated output must be balanced (even count). These three confirm AST extraction, LLM routing, and output formatting are all correct.


## Model source currency

- Anthropic model list: https://docs.anthropic.com/en/docs/about-claude/models/overview
- OpenAI model list: https://platform.openai.com/docs/models
- verifiedOn: 2026-08-07
- `LLM_MODEL` is required and has no implicit default; choose a supported model for the configured provider at setup time.