lucen-cli 0.1.0__py3-none-any.whl

lucen/progress.py

@@ -0,0 +1,96 @@
+"""Progress for a script: silent by default, the CLI's lines on request, or a callback.
+
+The CLI's engine reports through :class:`lucen_cli.progress.Reporter` —
+milestones (`step`) and byte counters (`bytes`) on stderr. A script wants one
+of three things, and :func:`make_reporter` builds the reporter for each:
+
+- **silent** (the default): nothing is written anywhere; the result object is
+  the whole answer,
+- ``progress=True``: exactly what `lucen push` prints — a rich bar on a
+  terminal, one plain line per decile in a log,
+- ``on_progress=callback``: every milestone and byte count as a
+  :class:`ProgressEvent`, for a notebook widget or a log line of your own.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Literal
+
+from lucen_cli.progress import ByteTask, Reporter
+
+
+@dataclass(frozen=True)
+class ProgressEvent:
+    """One thing the engine did.
+
+    `kind` is ``step`` (a milestone: `message`), ``bytes`` (a byte counter
+    `label` moved to `done` of `total`), or ``log`` (a line of a run's log
+    while it is followed: `message`, with `stream` ``stdout`` for the
+    trainer's own output and ``stderr`` for a notice the hub wrote).
+    """
+
+    kind: Literal["step", "bytes", "log"]
+    message: str = ""
+    label: str | None = None
+    done: int = 0
+    total: int = 0
+    stream: str | None = None
+
+    @property
+    def fraction(self) -> float | None:
+        """`done / total` for a byte event with a known total, else None."""
+        if self.kind != "bytes" or self.total <= 0:
+            return None
+        return min(1.0, self.done / self.total)
+
+
+ProgressCallback = Callable[[ProgressEvent], None]
+
+
+class _CallbackTask:
+    def __init__(self, callback: ProgressCallback, label: str, total: int) -> None:
+        self._callback = callback
+        self._label = label
+        self._total = total
+        self._done = 0
+        self._lock = threading.Lock()
+
+    def advance(self, count: int) -> None:
+        with self._lock:
+            self._done += count
+            done = self._done
+        self._callback(ProgressEvent("bytes", label=self._label, done=done, total=self._total))
+
+
+class CallbackReporter(Reporter):
+    """A :class:`Reporter` whose every milestone and byte count is one callback call."""
+
+    def __init__(self, callback: ProgressCallback) -> None:
+        super().__init__(quiet=True)
+        self._callback = callback
+
+    def step(self, message: str) -> None:
+        self._callback(ProgressEvent("step", message=message))
+
+    @contextmanager
+    def bytes(self, label: str, total: int) -> Iterator[ByteTask]:
+        task = _CallbackTask(self._callback, label, total)
+        self._callback(ProgressEvent("bytes", label=label, done=0, total=total))
+        yield task
+
+    def log(self, line: str, *, stream: str) -> None:
+        self._callback(ProgressEvent("log", message=line, stream=stream))
+
+
+def make_reporter(*, progress: bool, on_progress: ProgressCallback | None) -> Reporter:
+    """The reporter for one call: a callback wins, then the CLI's own lines, then silence."""
+    if on_progress is not None:
+        return CallbackReporter(on_progress)
+    return Reporter(quiet=not progress)
+
+
+__all__ = ["CallbackReporter", "ProgressCallback", "ProgressEvent", "make_reporter"]