diff --git a/src/cmos/cli.py b/src/cmos/cli.py index 3fdd2d8..3559304 100644 --- a/src/cmos/cli.py +++ b/src/cmos/cli.py @@ -9,6 +9,7 @@ from __future__ import annotations import argparse import sys +from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Callable @@ -17,16 +18,36 @@ from cmos.parser import split_bibliography FormatterFn = Callable[[str], str] +DEFAULT_CONCURRENCY = 8 -def reformat_draft(text: str, formatter: FormatterFn | None = None) -> str: + +def reformat_draft( + text: str, + formatter: FormatterFn | None = None, + concurrency: int = DEFAULT_CONCURRENCY, +) -> str: """Rewrite the bibliography section of ``text`` using ``formatter``. The formatter takes a single messy entry and returns the CMOS 18 form. If omitted, the real OpenAI-backed formatter is used. + + Entries are reformatted concurrently using a thread pool of size + ``concurrency``. OpenAI SDK calls are synchronous but network-bound, so + threads give real speedup. ``ThreadPoolExecutor.map`` preserves input + order regardless of which calls finish first, so the output is + deterministic. + + Set ``concurrency=1`` for sequential execution (useful for + deterministic debugging or when hitting rate limits). """ fmt = formatter or format_bibliography_entry parsed = split_bibliography(text) - rewritten = [fmt(entry) for entry in parsed.entries] + + if concurrency <= 1 or len(parsed.entries) <= 1: + rewritten = [fmt(entry) for entry in parsed.entries] + else: + with ThreadPoolExecutor(max_workers=concurrency) as pool: + rewritten = list(pool.map(fmt, parsed.entries)) pieces: list[str] = [] if parsed.before: @@ -46,11 +67,20 @@ def main(argv: list[str] | None = None) -> int: fmt = sub.add_parser("format", help="Reformat a draft's bibliography to CMOS 18.") fmt.add_argument("path", type=Path, help="Path to a markdown draft.") + fmt.add_argument( + "--concurrency", + type=int, + default=DEFAULT_CONCURRENCY, + help=( + "Number of parallel formatter calls (default: " + f"{DEFAULT_CONCURRENCY}). Set to 1 for serial execution." + ), + ) args = ap.parse_args(argv) if args.command == "format": text = args.path.read_text(encoding="utf-8") - sys.stdout.write(reformat_draft(text)) + sys.stdout.write(reformat_draft(text, concurrency=args.concurrency)) return 0 return 2 diff --git a/tests/test_cli.py b/tests/test_cli.py index ddb375c..725ae7a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -66,3 +66,38 @@ def test_reformat_draft_bibliography_heading_preserved(): draft = "## Bibliography\n\nentry.\n" output = reformat_draft(draft, formatter=_fake_formatter) assert "## Bibliography" in output + + +def test_reformat_draft_preserves_order_under_concurrency(): + # With concurrent execution the formatter is called on all entries in + # parallel; the CLI must reassemble them in input order regardless of + # which call finishes first. This test uses a fake formatter that + # sleeps based on the entry content so earlier entries finish AFTER + # later ones if order were naively tied to completion. + import time + + def slow_fake(messy: str) -> str: + # Entries with lower index sleep longer so they finish last. + n = int(messy.split()[-1]) + time.sleep(0.05 * (5 - n)) + return f"FORMATTED({messy})" + + draft = """\ +## Bibliography + +entry 0 +entry 1 +entry 2 +entry 3 +entry 4 +""" + output = reformat_draft(draft, formatter=slow_fake, concurrency=4) + # Entries must appear in input order. + lines = [l for l in output.splitlines() if l.startswith("FORMATTED(")] + assert lines == [ + "FORMATTED(entry 0)", + "FORMATTED(entry 1)", + "FORMATTED(entry 2)", + "FORMATTED(entry 3)", + "FORMATTED(entry 4)", + ]