Writing documentation is the task developers deprioritize and teams regret. Legacy codebases accumulate thousands of undocumented functions, and onboarding becomes weeks of archaeology. Language models now make it feasible to generate accurate, contextual docstrings directly from source — and wiring this into a working tool takes less than an hour.
The Architecture
The pipeline is: parse → filter → prompt → patch. Each step is testable in isolation. No frameworks, no magic — just Python's ast module, direct API calls, and careful line-number arithmetic.
The core design decision: work with AST nodes and line numbers, not regex. Regex breaks on decorators, multi-line signatures, and nested functions. The AST does not.
Parsing Functions from Source
Python's ast module gives us precise control. We need the function source text, its line number for insertion, and whether it already has a docstring — that last bit ensures idempotency when you run the tool on a partially-documented codebase.
import ast
from dataclasses import dataclass
@dataclass
class FunctionInfo:
name: str
source: str
lineno: int
col_offset: int
has_docstring: bool
def extract_functions(source_code: str) -> list[FunctionInfo]:
tree = ast.parse(source_code)
results = []
for node in ast.walk(tree):
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
has_doc = (
node.body
and isinstance(node.body[0], ast.Expr)
and isinstance(node.body[0].value, ast.Constant)
and isinstance(node.body[0].value.value, str)
)
func_source = ast.get_source_segment(source_code, node)
if func_source:
results.append(FunctionInfo(
name=node.name,
source=func_source,
lineno=node.lineno,
col_offset=node.col_offset,
has_docstring=has_doc,
))
return results
ast.get_source_segment requires Python 3.8+. We track col_offset so the inserted docstring gets the correct indentation — class methods are indented 4 spaces deeper than module-level functions, and getting this wrong produces a SyntaxError on the next parse.
Generating Docstrings with a Language Model
Prompt structure matters more than model choice. A structured prompt that specifies Google-style format — one-line summary, Args, Returns, Raises — produces output you'd actually commit. A vague "document this function" produces noise.
import httpx
import os
def generate_docstring(func: FunctionInfo) -> str:
prompt = (
"Generate a concise Google-style docstring for this Python function.\n"
"Include: one-line summary, Args (if any), Returns, Raises (if applicable).\n"
"Output only the docstring body — no triple quotes, no function signature.\n\n"
f"Function:\n{func.source}"
)
resp = httpx.post(
os.environ["LLM_API_BASE"] + "/chat/completions",
headers={"Authorization": f"Bearer {os.environ['LLM_API_KEY']}"},
json={
"model": os.environ.get("LLM_MODEL", "gpt-4o-mini"),
"messages": [{"role": "user", "content": prompt}],
"temperature": 0.2,
"max_tokens": 256,
},
timeout=30,
)
resp.raise_for_status()
return resp.json()["choices"][0]["message"]["content"].strip()
Temperature at 0.2 is deliberate — you want factual, repeatable output, not creative reinterpretation. The same function should produce near-identical docstrings on consecutive runs. Any OpenAI-compatible endpoint works here: a local inference server, a managed API, or a fine-tuned model. Swap LLM_API_BASE and you are done.
Patching the Source Without Breaking It
String replacement fails on duplicate function names in different classes. The correct approach: insert docstrings by line number, processing functions in reverse order.
def patch_source(
source_code: str,
functions: list[FunctionInfo],
docstrings: dict[str, str],
) -> str:
lines = source_code.splitlines(keepends=True)
# Reverse order so earlier insertions don't shift later line numbers
to_patch = sorted(
[(f, docstrings[f.name]) for f in functions if f.name in docstrings],
key=lambda x: x[0].lineno,
reverse=True,
)
for func, content in to_patch:
# Walk forward past multi-line signatures to find the closing colon
insert_at = func.lineno
while insert_at <= len(lines):
if lines[insert_at - 1].rstrip().endswith(":"):
break
insert_at += 1
indent = " " * (func.col_offset + 4)
doc_lines = content.splitlines()
opening = indent + '"""' + doc_lines[0] + "\n"
middle = "".join(indent + line + "\n" for line in doc_lines[1:])
closing = indent + '"""' + "\n"
lines.insert(insert_at, opening + middle + closing)
return "".join(lines)
The reverse-order insight: insert a 5-line docstring at line 20 and every function below shifts by 5. Process top-to-bottom and you corrupt every subsequent line number. Process bottom-to-top and the problem disappears.
The "walk forward to find the colon" loop handles multi-line function signatures — increasingly common with type annotations and keyword arguments that span several lines.
Wiring It Into a CLI
import argparse
def main():
parser = argparse.ArgumentParser(description="Generate docstrings from source")
parser.add_argument("file", help="Python file to document")
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
with open(args.file) as f:
source = f.read()
funcs = [fn for fn in extract_functions(source) if not fn.has_docstring]
if not funcs:
print("No undocumented functions found.")
return
print(f"Generating {len(funcs)} docstring(s)...")
docstrings = {}
for fn in funcs:
print(f" → {fn.name}")
docstrings[fn.name] = generate_docstring(fn)
patched = patch_source(source, funcs, docstrings)
if args.dry_run:
print(patched)
else:
with open(args.file, "w") as f:
f.write(patched)
print(f"Done. {len(docstrings)} docstring(s) written.")
if __name__ == "__main__":
main()
This drops cleanly into a pre-commit hook: fail the build if any undocumented functions exist, or run in auto-fix mode to generate them before the commit lands.
Security Before You Deploy This Internally
Before rolling this out to a team, audit what leaves your network. Source code often carries internal service names, database schemas, API structures, and proprietary business logic. Two concrete mitigations worth implementing:
Self-host the model — the architecture above works with any OpenAI-compatible endpoint. Point LLM_API_BASE at a local inference server and no code leaves your infrastructure. For teams evaluating their LLM deployment security posture, we publish free security hardening checklists that include AI inference infrastructure controls covering data exposure, model access policies, and audit logging.
Strip before sending — replace actual identifiers with generic placeholders before the prompt and restore them after. An AST-based approach makes this mechanical: rename all ast.Name nodes to varN, generate the docstring, then restore the originals in a post-processing step.
Developer tooling that silently exfiltrates source code to third-party APIs is a documented audit finding in regulated environments. The architecture above keeps that control explicit.
The Takeaway
The LLM call is the easy part. The AST manipulation and backwards patching logic deserve thorough unit tests — write cases for nested functions, async generators, class methods, decorated functions with multi-line signatures, and files that mix documented and undocumented functions.
The model layer is intentionally thin: one function, two environment variables. Today it is a cloud API; next quarter it might be a fine-tuned model running on-premises. Keep it swappable and the rest of the tool stays useful regardless of what the ecosystem looks like in six months.
I run AYI NEDJIMI Consultants, a cybersecurity consulting firm. We publish free security hardening checklists — PDF and Excel.
Top comments (0)