arborist/arborist/progress.py

"""Periodic stderr progress reporting for long-running ingests/distillation.

Stdlib only. Prints on a rate-limited schedule (default every 2 seconds) so
ingests of 100k+ docs give live feedback without spamming. Optional total
estimate produces a percent + ETA.

Multiple parallel shard processes will interleave their lines; each can be
given a `prefix` to disambiguate.
"""

from __future__ import annotations

import sys
import time
from typing import IO


class Progress:
    """Rate-limited progress reporter to stderr (or any stream)."""

    def __init__(
        self,
        *,
        interval: float = 2.0,
        total_estimate: int | None = None,
        prefix: str = "",
        stream: IO[str] = sys.stderr,
    ):
        self.start = time.monotonic()
        self.last_print = self.start
        self.last_seen = 0
        self.interval = interval
        self.total_estimate = total_estimate
        self.prefix = prefix
        self.stream = stream

    def tick(
        self,
        seen: int,
        *,
        inserted: int | None = None,
        force: bool = False,
    ) -> None:
        """Maybe print a status line. Call after each batch flush."""
        now = time.monotonic()
        if not force and (now - self.last_print) < self.interval:
            return
        elapsed = now - self.start
        delta_seen = max(0, seen - self.last_seen)
        delta_t = max(now - self.last_print, 1e-6)
        rate_now = delta_seen / delta_t
        rate_avg = seen / elapsed if elapsed > 0 else 0.0

        bits = [
            f"{self.prefix}[{_format_elapsed(elapsed)}]",
            f"{seen:>10,} seen",
        ]
        if inserted is not None:
            bits.append(f"{inserted:>9,} new")
        bits.append(f"{rate_now:>5.0f} now")
        bits.append(f"({rate_avg:>4.0f} avg) docs/s")
        if self.total_estimate and rate_now > 0:
            pct = 100.0 * min(1.0, seen / self.total_estimate)
            remaining = max(0, self.total_estimate - seen)
            eta = remaining / rate_now
            bits.append(f"{pct:>5.1f}%")
            bits.append(f"ETA {_format_elapsed(eta)}")

        print("  " + " | ".join(bits), file=self.stream, flush=True)
        self.last_print = now
        self.last_seen = seen

    def done(self, seen: int, *, inserted: int | None = None) -> None:
        """Print a final status line regardless of interval."""
        self.tick(seen, inserted=inserted, force=True)


def _format_elapsed(secs: float) -> str:
    secs = int(secs)
    h, rem = divmod(secs, 3600)
    m, s = divmod(rem, 60)
    if h:
        return f"{h}h{m:02d}m"
    if m:
        return f"{m:>2}m{s:02d}s"
    return f"   {s:>2}s"