step

@dataclass(frozen=True, slots=True)
class Reading:
    """A step's count and note at one instant.

    Attributes:
        done: How many units the step has advanced through.
        total: How many it will, if known.
        message: The last ``note``, if any.
        fields: The named values the notes carried, latest per name.
    """

    done: int
    total: int | None
    message: str | None
    fields: dict[str, JSONValue]

class Step:
    """One span of work; ``step()`` opens one and the methods below advance it.

    Attributes:
        id: A fresh identifier.
        parent_id: The ambient step's ``id`` when this one opened, if any.
        label: What the step is doing, as a page or terminal shows it.
        identity: The function identity a memoized call reports, else ``None``.
        call_key: The call key a memoized call reports, else ``None``.
        workspace: The workspace a memoized call computes in, else ``None``.
        scratch: The scratch the step reports to, if one was active when it opened.
        started_at: When the step opened, as ``time.time()``.
    """

    __slots__ = (
        "_dirty",
        "_done",
        "_fields",
        "_lock",
        "_message",
        "call_key",
        "id",
        "identity",
        "label",
        "parent_id",
        "scratch",
        "started_at",
        "total",
        "workspace",
    )

    def __init__(
        self,
        label: str,
        total: int | None = None,
        *,
        identity: str | None = None,
        call_key: CallKey | None = None,
        workspace: Path | None = None,
    ) -> None:
        """A step not yet open, beneath the ambient step and reporting to the active scratch."""
        parent = active_step.get()
        self.id = uuid.uuid4().hex
        self.parent_id = parent.id if parent is not None else None
        self.label = label
        self.total = total
        self.identity = identity
        self.call_key = call_key
        self.workspace = workspace
        self.scratch: Scratch | None = current_scratch()
        self.started_at = time.time()
        self._done = 0
        self._message: str | None = None
        self._fields: dict[str, JSONValue] = {}
        self._dirty = True
        self._lock = threading.Lock()

    def advance(self, count: int = 1) -> None:
        """Count ``count`` more units done."""
        with self._lock:
            self._done += count
            self._dirty = True

    def track[T](self, iterable: Iterable[T]) -> Iterator[T]:
        """``iterable``, advancing this step by one as each element is consumed.

        ```python
        for batch in training.track(batches):
            ...
        ```
        """
        for item in iterable:
            yield item
            self.advance()

    def note(self, message: str | None = None, **fields: JSONValue) -> None:
        """Say what the step is doing: the last message wins, ``fields`` update by name.

        ```python
        training.note(f"loss {loss:.3f}", loss=loss, psnr=psnr)
        ```
        """
        with self._lock:
            if message is not None:
                self._message = message
            self._fields.update(fields)
            self._dirty = True

    @property
    def elapsed(self) -> float:
        """Seconds since the step opened."""
        return time.time() - self.started_at

    def read(self, *, settle: bool = False) -> Reading:
        """The count and note right now; ``settle`` also clears the dirty mark."""
        with self._lock:
            if settle:
                self._dirty = False
            return Reading(self._done, self.total, self._message, dict(self._fields))

    @property
    def dirty(self) -> bool:
        """Whether the step has changed since the last settling ``read``."""
        return self._dirty

    def __repr__(self) -> str:
        """``Step(<label>)``."""
        return f"Step({self.label!r})"