#!/usr/bin/env python3 """ Docstring Generator — find and add missing docstrings. Scans Python files for functions/async functions lacking docstrings. Generates Google-style docstrings from function signature and body. Inserts them in place. Usage: python3 docstring_generator.py scripts/ # Fix in place python3 docstring_generator.py --dry-run scripts/ # Preview changes python3 docstring_generator.py --json scripts/ # Machine-readable output python3 docstring_generator.py path/to/file.py """ import argparse import ast import json import os import sys from pathlib import Path from typing import Optional, Tuple, List # --- Helper: turn snake_case into Title Case phrase --- def name_to_title(name: str) -> str: """Convert snake_case function name to a Title Case description.""" words = name.replace('_', ' ').split() if not words: return '' titled = [] for w in words: if len(w) <= 2: titled.append(w.upper()) else: titled.append(w[0].upper() + w[1:]) return ' '.join(titled) # --- Helper: extract first meaningful statement from body for summary --- def extract_body_hint(body: list[ast.stmt]) -> Optional[str]: """Look for an assignment or return that hints at function purpose.""" for stmt in body: if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant): continue # skip existing docstring placeholder # Assignment to a result-like variable? if isinstance(stmt, ast.Assign): for target in stmt.targets: if isinstance(target, ast.Name): var_name = target.id if var_name in ('result', 'msg', 'output', 'retval', 'value', 'response', 'data'): val = ast.unparse(stmt.value).strip() if val: return f"Compute or return {val}" # Return statement if isinstance(stmt, ast.Return) and stmt.value: ret = ast.unparse(stmt.value).strip() if ret: return f"Return {ret}" break return None # --- Generate a docstring string for a function --- def generate_docstring(func_node: ast.FunctionDef | ast.AsyncFunctionDef) -> str: """Build a Google-style docstring for the given function node.""" parts: list[str] = [] # Summary line summary = name_to_title(func_node.name) body_hint = extract_body_hint(func_node.body) if body_hint: summary = f"{summary}. {body_hint}" parts.append(summary) # Args section if there are parameters (excluding self/cls) args = func_node.args.args if args: arg_lines = [] for arg in args: if arg.arg in ('self', 'cls'): continue type_ann = ast.unparse(arg.annotation) if arg.annotation else 'Any' arg_lines.append(f"{arg.arg} ({type_ann}): Parameter {arg.arg}") if arg_lines: parts.append("\nArgs:\n " + "\n ".join(arg_lines)) # Returns section if func_node.returns: ret_type = ast.unparse(func_node.returns) parts.append(f"\nReturns:\n {ret_type}: Return value") elif any(isinstance(s, ast.Return) and s.value is not None for s in ast.walk(func_node)): parts.append("\nReturns:\n Return value") return '"""' + '\n'.join(parts) + '\n"""' # --- Transform source AST --- def process_source(source: str, filename: str) -> Tuple[str, List[str]]: """Add docstrings to all undocumented functions. Returns (new_source, [func_names]).""" try: tree = ast.parse(source) except SyntaxError as e: print(f" WARNING: Could not parse {filename}: {e}", file=sys.stderr) return source, [] class DocstringInserter(ast.NodeTransformer): def __init__(self): self.modified_funcs: list[str] = [] def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.FunctionDef: return self._process(node) def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AsyncFunctionDef: return self._process(node) def _process(self, node): existing_doc = ast.get_docstring(node) if existing_doc is not None: return node docstring_text = generate_docstring(node) doc_node = ast.Expr(value=ast.Constant(value=docstring_text)) node.body.insert(0, doc_node) ast.fix_missing_locations(node) self.modified_funcs.append(node.name) return node inserter = DocstringInserter() new_tree = inserter.visit(tree) if inserter.modified_funcs: return ast.unparse(new_tree), inserter.modified_funcs return source, [] # --- File discovery --- def iter_python_files(paths: list[str]) -> list[Path]: """Collect all .py files from provided paths.""" files: set[Path] = set() for p in paths: path = Path(p) if not path.exists(): print(f"WARNING: Path not found: {p}", file=sys.stderr) continue if path.is_file() and path.suffix == '.py': files.add(path.resolve()) elif path.is_dir(): for child in path.rglob('*.py'): if '.git' in child.parts or '__pycache__' in child.parts: continue files.add(child.resolve()) return sorted(files) def main(): parser = argparse.ArgumentParser(description="Generate docstrings for functions missing them") parser.add_argument('paths', nargs='+', help='Python files or directories to process') parser.add_argument('--dry-run', action='store_true', help='Show what would change without writing') parser.add_argument('--json', action='store_true', help='Output machine-readable JSON summary') parser.add_argument('-v', '--verbose', action='store_true', help='Print each file processed') args = parser.parse_args() files = iter_python_files(args.paths) if not files: print("No Python files found to process", file=sys.stderr) sys.exit(1) results = [] total_funcs = 0 for pyfile in files: try: original = pyfile.read_text(encoding='utf-8') except Exception as e: print(f" ERROR reading {pyfile}: {e}", file=sys.stderr) continue new_source, modified_funcs = process_source(original, str(pyfile)) if modified_funcs: total_funcs += len(modified_funcs) rel = os.path.relpath(pyfile) if args.verbose: print(f" {rel}: +{len(modified_funcs)} docstrings") results.append({'file': str(pyfile), 'functions': modified_funcs}) if not args.dry_run: pyfile.write_text(new_source, encoding='utf-8') elif args.verbose: print(f" {rel}: no changes") if args.json: summary = {'total_files_modified': len(results), 'total_functions': total_funcs, 'files': results} print(json.dumps(summary, indent=2)) else: print(f"Generated docstrings for {total_funcs} functions across {len(results)} files") if args.dry_run: print(" (dry run — no files written)") return 0 if __name__ == '__main__': sys.exit(main())