compounding-intelligence/scripts/docstring_generator.py

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