#!/usr/bin/env python3 """ Apply Docstrings from suggestions JSON to source files. Reads /tmp/documentation-suggestions.json (generated by generate-docstrings.py), inserts Google-style docstrings into source files using AST-based function lookup. Safety: validates each file with py_compile after modification; reverts on failure. """ import ast import json import py_compile import sys from pathlib import Path PROTECTED_PATTERNS = [ ".github/workflows/", ".env", "ods/installers/", "ods/ods-cli", "ods/config/", ] def is_protected(file_path: str) -> bool: """Check if a file path matches any protected patterns.""" for pattern in PROTECTED_PATTERNS: if file_path.startswith(pattern) or f"/{pattern}" in file_path: return True return False def find_function_info(source: str, function_name: str) -> dict | None: """Use AST to find function line and check if it already has a docstring.""" try: tree = ast.parse(source) except SyntaxError: return None for node in ast.walk(tree): if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): if node.name == function_name: has_docstring = ast.get_docstring(node) is not None return { "lineno": node.lineno, "has_docstring": has_docstring, "body_start": node.body[0].lineno if node.body else node.lineno + 1, } return None def find_def_end_line(lines: list[str], start_idx: int) -> int: """Find the 0-based index of the last line of a def statement.""" depth = 0 for i in range(start_idx, len(lines)): line = lines[i] depth += line.count("(") - line.count(")") if depth <= 0 and ":" in line: stripped = line.rstrip() if stripped.endswith(":"): return i colon_pos = stripped.rfind(":") if colon_pos >= 0: return i return start_idx def get_body_indent(lines: list[str], def_end_idx: int) -> str: """Determine the indentation level of the function body.""" for i in range(def_end_idx + 1, min(def_end_idx + 5, len(lines))): line = lines[i] stripped = line.strip() if stripped and not stripped.startswith("#"): return line[: len(line) - len(line.lstrip())] def_line = lines[def_end_idx] if def_end_idx < len(lines) else "" def_indent = def_line[: len(def_line) - len(def_line.lstrip())] return def_indent + " " def format_docstring(docstring: str, indent: str) -> list[str]: """Format a docstring with proper indentation as lines to insert.""" doc_lines = docstring.split("\n") if len(doc_lines) == 1: return [f'{indent}"""{doc_lines[0]}"""\n'] result = [f'{indent}"""{doc_lines[0]}\n'] for line in doc_lines[1:]: if line.strip(): result.append(f"{indent}{line}\n") else: result.append("\n") result.append(f'{indent}"""\n') return result def apply_docstrings(suggestions_path: str) -> dict: """Apply generated docstring suggestions to Python source files.""" with open(suggestions_path, "r") as f: data = json.load(f) functions = data.get("functions_documented", []) if not functions: print("No docstring suggestions to apply.") return { "files_modified": 0, "docstrings_inserted": 0, "files_reverted": 0, "skipped_existing": 0, } by_file: dict[str, list] = {} for func in functions: file_path = func.get("file", "") if not file_path: continue if is_protected(file_path): print(f" Skipping protected file: {file_path}") continue if not Path(file_path).exists(): print(f" Skipping missing file: {file_path}") continue by_file.setdefault(file_path, []).append(func) files_modified = 0 docstrings_inserted = 0 files_reverted = 0 skipped_existing = 0 for file_path, file_funcs in by_file.items(): print(f"\nProcessing {file_path} ({len(file_funcs)} functions)...") original_content = Path(file_path).read_text() source = original_content lines = source.splitlines(keepends=True) located = [] for func in file_funcs: info = find_function_info(source, func["function"]) if info is None: print( f" Could not find function '{func['function']}' in AST, skipping" ) continue if info["has_docstring"]: print( f" Function '{func['function']}' already has docstring, skipping" ) skipped_existing += 1 continue located.append((info["lineno"], func)) located.sort(key=lambda x: x[0], reverse=True) inserted_in_file = 0 for actual_line, func in located: docstring = func.get("docstring", "").strip() if not docstring: continue start_idx = actual_line - 1 if start_idx >= len(lines) or start_idx < 0: continue def_end_idx = find_def_end_line(lines, start_idx) body_indent = get_body_indent(lines, def_end_idx) doc_lines = format_docstring(docstring, body_indent) insert_point = def_end_idx + 1 lines[insert_point:insert_point] = doc_lines inserted_in_file += 1 print( f" Inserted docstring: {func['function']} (after line {def_end_idx + 1})" ) if inserted_in_file == 0: continue new_source = "".join(lines) Path(file_path).write_text(new_source) try: py_compile.compile(file_path, doraise=True) files_modified += 1 docstrings_inserted += inserted_in_file print(f" Validated: {file_path} ({inserted_in_file} docstrings)") except py_compile.PyCompileError as e: print(f" REVERT: {file_path} failed compilation: {e}") Path(file_path).write_text(original_content) files_reverted += 1 return { "files_modified": files_modified, "docstrings_inserted": docstrings_inserted, "files_reverted": files_reverted, "skipped_existing": skipped_existing, } def main() -> None: """Load and apply generated docstrings.""" suggestions_path = ( sys.argv[1] if len(sys.argv) > 1 else "/tmp/documentation-suggestions.json" ) if not Path(suggestions_path).exists(): print(f"Suggestions file not found: {suggestions_path}") print("No docstrings to apply.") return print(f"Loading suggestions from {suggestions_path}...") result = apply_docstrings(suggestions_path) print("\n## Documentation Application Summary\n") print(f"- **Files modified**: {result['files_modified']}") print(f"- **Docstrings inserted**: {result['docstrings_inserted']}") print(f"- **Skipped** (already had docstring): {result['skipped_existing']}") print(f"- **Files reverted** (compilation failed): {result['files_reverted']}") if __name__ == "__main__": main()