light-heart-labs--dreamserver/.github/scripts/apply-docstrings.py

#!/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()