--- title: "Code Documentation Generator" description: "A shippable internal tool that directly reduces team toil — the kind of project that justifies 'AI tooling' on a senior engineer's resume with a real..." source: "https://confidentprep.com/paths/ai-augmented-engineering/code-doc-generator/" path: "AI-Augmented Engineering" repository: "https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/code-doc-generator" token_estimate: 3988 --- # Code Documentation Generator ## Overview Parses Python files with the `ast` module to extract function signatures and type hints, then generates LLM-written README documentation for every public function — with a diff mode that previews exactly what changed before anything is written. The diff-preview pattern matters: it's what makes an AI tool safe to point at a real codebase instead of something you only trust on toy examples. ## What to do | Field | Value | |---|---| | Path | 2 — AI-Augmented Engineering | | Position | 4 of 12 | | Difficulty | 🟢 Local only | | Estimated time | 3 hours | | AWS cost | None | --- ### Agent Pickup Instructions ```bash # 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 ` 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: [--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 ```bash # 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)** ```python @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 `# ` 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 ` — 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 ` 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 ` 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 ```markdown # 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 ```bash 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: ```markdown ### 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`: ```python 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 ." and exit 0 | | `--diff` used without `--output` | `main.py` | Print "Error: --diff requires --output " 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. ## Source code A reference implementation of Code Documentation Generator — README, guide, source and tests — is at https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/code-doc-generator.