sparse.py

import argparse
import concurrent.futures
import os
import signal
import sys
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

from outputter.writer import merge_results, write_json
from parser.md_file import parse_md_file

BAR_WIDTH = 30
_executor = None

# Default input directories per --mode, anchored to the script location so
# the defaults work even when sparse.py is invoked from a different cwd.
_SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_INPUT_DIR = {
    "sdk": _SCRIPT_DIR / "sdk-api" / "sdk-api-src" / "content",
    "driver": _SCRIPT_DIR / "windows-driver-docs-ddi" / "wdk-ddi-src" / "content",
}


def _sigint_handler(signum, frame):
    procs = list(_executor._processes.values()) if _executor is not None else []
    sys.stderr.write(f"\n\n  \u26a0  Interrupted — killing {len(procs)} worker(s)...")
    sys.stderr.flush()
    for proc in procs:
        try:
            proc.terminate()
        except Exception:
            pass
    sys.stderr.write(" done.\n")
    sys.stderr.flush()
    os._exit(130)


def process_chunk(file_paths, mode):
    """Process a chunk of md files. Returns dict keyed by function name."""
    results = {}
    for fp in file_paths:
        parsed = parse_md_file(fp, mode=mode)
        if parsed is None:
            continue
        func_name, data = parsed
        if func_name:
            results[func_name] = data
    return results


def chunk_list(lst, size):
    """Split list into chunks of given size."""
    for i in range(0, len(lst), size):
        yield lst[i : i + size]


def print_progress(done, total, entries, elapsed, failed=0):
    """Print a progress bar to stderr."""
    frac = done / total if total else 1
    filled = int(BAR_WIDTH * frac)
    bar = "\u2588" * filled + "\u2591" * (BAR_WIDTH - filled)
    pct = frac * 100
    fail_str = f"  \u2717 {failed} failed" if failed else ""
    sys.stderr.write(
        f"\r  [{bar}] {pct:5.1f}%  "
        f"{done}/{total} chunks  "
        f"{entries:,} entries  "
        f"{elapsed:.1f}s{fail_str}"
    )
    sys.stderr.flush()


def main():
    # NOTE(gabriela): so we can ctrl+c safely at any time, killing all 
    # child processes.
    signal.signal(signal.SIGINT, _sigint_handler)

    # ───────────────────────────────── Arguments ──────────────────────────────── #

    parser = argparse.ArgumentParser(
        description="Parse MSDN Native Function .md files to JSON (sdk-api or windows-driver-docs-ddi)"
    )
    parser.add_argument(
        "input_dir",
        nargs="?",
        default=None,
        help=(
            "Root content directory to search for nf-*.md files. "
            "Defaults to the matching submodule for --mode "
            "(sdk-api/sdk-api-src/content for --mode sdk, "
            "windows-driver-docs-ddi/wdk-ddi-src/content for --mode driver)."
        ),
    )
    parser.add_argument(
        "--mode",
        choices=["sdk", "driver"],
        default="sdk",
        help="Frontmatter dialect: 'sdk' for sdk-api (default), 'driver' for windows-driver-docs-ddi",
    )
    parser.add_argument(
        "-o", "--output", default="output.json", help="Output JSON file path"
    )
    # NOTE(gabriela): From testing, it seems like 32/64 is around the best
    # split of data in chunks relative to performance.
    parser.add_argument(
        "--chunk-size", type=int, default=64, help="Files per worker chunk"
    )
    # NOTE(gabriela): 32 seems to be the best as it doesn't create too many
    # processes, meaning there isn't a long delay in the start of the app,
    # and it scales with our default chunk size really nicely.
    parser.add_argument(
        "--workers",
        default=32,
        help='Number of parallel workers, or "max" to use _MAX_WINDOWS_WORKERS',
    )
    parser.add_argument("--silent", action="store_true", help="Suppress all output")

    # ──────────────────────────────────────────────────────────────────────────── #

    args = parser.parse_args()

    # NOTE(gabriela): offer a way to use max workers count
    if args.workers == "max":
        args.workers = concurrent.futures.process._MAX_WINDOWS_WORKERS
    else:
        args.workers = int(args.workers)


    log = (lambda *a, **kw: None) if args.silent else lambda *a, **kw: print(*a, **kw)

    if args.input_dir is None:
        args.input_dir = str(DEFAULT_INPUT_DIR[args.mode])

    input_path = Path(args.input_dir)
    if not input_path.is_dir():
        print(f"Error: {args.input_dir} is not a directory", file=sys.stderr)
        if str(DEFAULT_INPUT_DIR[args.mode]) == args.input_dir:
            print(
                "Hint: did you run `git submodule update --init`?",
                file=sys.stderr,
            )
        sys.exit(1)

    # Validate output path BEFORE doing any parsing — fail fast on bad dirs.
    output_path = Path(args.output)
    out_parent = output_path.parent if str(output_path.parent) else Path(".")
    if not out_parent.exists():
        print(
            f"Error: output directory {out_parent} does not exist",
            file=sys.stderr,
        )
        sys.exit(1)
    if not out_parent.is_dir():
        print(f"Error: {out_parent} is not a directory", file=sys.stderr)
        sys.exit(1)
    if output_path.exists() and output_path.is_dir():
        print(f"Error: {output_path} is a directory, not a file", file=sys.stderr)
        sys.exit(1)

    # Glob all nf-*.md files
    all_files = sorted(str(p) for p in input_path.rglob("nf-*.md"))
    log(f"  Found {len(all_files):,} nf-*.md files")

    if not all_files:
        print("No files found, exiting.", file=sys.stderr)
        sys.exit(1)

    # Process in parallel chunks
    chunks = list(chunk_list(all_files, args.chunk_size))
    log(f"  Processing with {args.workers} workers...\n")

    # Slot-indexed so the merge order is deterministic regardless of
    # which worker finishes first. Required: two source files can produce
    # the same function name (A/W pairs, COM overloads); the last-merged
    # one wins, so non-deterministic merge order produces different output
    # across runs.
    chunk_results = [None] * len(chunks)
    done = 0
    total_entries = 0
    failed = 0
    t0 = time.time()

    global _executor
    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        _executor = executor
        futures = {
            executor.submit(process_chunk, chunk, args.mode): i
            for i, chunk in enumerate(chunks)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                result = future.result()
                chunk_results[idx] = result
                total_entries += len(result)
            except Exception:
                failed += 1
            done += 1
            if not args.silent:
                print_progress(
                    done, len(chunks), total_entries, time.time() - t0, failed
                )

    if not args.silent:
        sys.stderr.write("\n\n")
        sys.stderr.flush()

    # Merge and write
    merged = merge_results(chunk_results)
    write_json(merged, args.output)
    log(f"  {len(merged):,} entries -> {args.output}")


if __name__ == "__main__":
    main()