watlowlib.streaming¶

@asynccontextmanager
async def record(
    source: PollSource,
    *,
    parameters: Sequence[str | int],
    rate_hz: float,
    duration: float | None = None,
    names: Sequence[str] | None = None,
    instances: Sequence[int] = (1,),
    overflow: OverflowPolicy = OverflowPolicy.BLOCK,
    buffer_size: int = 64,
    auto_reconnect: bool = False,
    reconnect_factory: Callable[[], Awaitable[PollSource]] | None = None,
) -> AsyncGenerator[Recording[Sequence[Sample]]]:
    """Record polled samples into a receive stream at an absolute cadence.

    Usage::

        async with record(
            controller, parameters=["process_value", "setpoint"], rate_hz=2, duration=10
        ) as recording:
            async for batch in recording.stream:
                for sample in batch:
                    print(sample.parameter, sample.value)
            # recording.summary is live; recording.summary.finished_at is None
            # while running and set on CM exit.

    The CM yields a :class:`Recording[Sequence[Sample]]` exposing
    ``.stream`` (async iterator of per-tick :class:`Sample` batches),
    ``.summary`` (live :class:`AcquisitionSummary` — recorder is sole
    writer), and ``.rate_hz`` (the cadence the recorder is running
    at).

    Each batch is a flat :class:`Sequence` — one entry per (device,
    parameter, instance) read that succeeded. Failed reads are dropped
    by the source and logged at WARN.

    Args:
        source: Any :class:`PollSource` (a :class:`Controller` or a
            :class:`WatlowManager`).
        parameters: Parameter names or registry IDs to poll each tick.
        rate_hz: Target cadence. Absolute targets are computed
            ``target[n] = start + n * (1 / rate_hz)``. Must be > 0.
        duration: Total acquisition duration in seconds. ``None``
            means "until the caller exits the CM".
        names: Subset of device names to poll per tick. ``None`` polls
            everything the source manages. Ignored for solo controllers.
        instances: 1-indexed loop / channel numbers per device. Single-
            loop devices use ``(1,)``.
        overflow: Backpressure policy when the receive-stream buffer
            is full. See :class:`OverflowPolicy`.
        buffer_size: Receive-stream capacity, in per-tick batches.
        auto_reconnect: When ``True``, treat
            :class:`WatlowConnectionError` raised by ``source.poll_many``
            as a transient transport drop rather than a fatal error.
            The producer logs ``recorder.disconnected``, waits per the
            backoff schedule, and either rebuilds the source via
            ``reconnect_factory`` (if supplied) or simply retries the
            same ``source.poll_many`` on the next tick. ``samples_late``
            ticks up for each tick missed during the gap.
        reconnect_factory: When supplied alongside ``auto_reconnect``,
            invoked to rebuild the :class:`PollSource` after a
            disconnect. Useful when the source's transport needs to be
            re-opened explicitly (e.g. a fresh
            :func:`watlowlib.open_device` call). The returned source
            replaces ``source`` for subsequent ticks. Without a
            factory, the recorder relies on ``source.poll_many`` itself to
            recover (which works for callers that wrap their own
            transport-reopen logic inside ``poll_many``).

    Yields:
        A :class:`Recording[Sequence[Sample]]` exposing ``.stream``,
        ``.summary``, and ``.rate_hz``.

    Raises:
        ValueError: ``rate_hz <= 0``, ``duration <= 0``, or
            ``buffer_size < 1``.
    """
    if rate_hz <= 0:
        raise ValueError(f"rate_hz must be > 0, got {rate_hz!r}")
    if duration is not None and duration <= 0:
        raise ValueError(f"duration must be > 0 or None, got {duration!r}")
    if buffer_size < 1:
        raise ValueError(f"buffer_size must be >= 1, got {buffer_size!r}")
    if not parameters:
        raise ValueError("parameters must be a non-empty sequence")

    period = 1.0 / rate_hz
    total_ticks = None if duration is None else max(1, round(duration * rate_hz))

    send_stream, receive_stream = anyio.create_memory_object_stream[Sequence[Sample]](
        max_buffer_size=buffer_size,
    )
    # Producer-side clone of the receive stream — used to evict the
    # oldest queued batch under DROP_OLDEST. Cloning here (before the
    # consumer starts iterating) keeps the eviction path off the
    # consumer's iterator and avoids racing with it.
    drop_rx = receive_stream.clone()

    started_at = datetime.now(UTC)
    summary = AcquisitionSummary(started_at=started_at)
    tick_durations_ms: list[float] = []
    _logger.info(
        "recorder.start rate_hz=%s duration_s=%s overflow=%s buffer_size=%s names=%s",
        rate_hz,
        duration,
        overflow.value,
        buffer_size,
        list(names) if names is not None else None,
    )

    async with anyio.create_task_group() as tg, receive_stream:

        async def _producer_entrypoint() -> None:
            await _run_producer(
                source,
                send_stream,
                drop_rx,
                tuple(parameters),
                tuple(instances),
                names,
                period,
                total_ticks,
                overflow,
                summary,
                tick_durations_ms,
                auto_reconnect=auto_reconnect,
                reconnect_factory=reconnect_factory,
            )

        # ``start_soon`` returns a TaskHandle (anyio >= 4.14); the producer is
        # fire-and-forget and cancelled via the group below, so discard it.
        _ = tg.start_soon(_producer_entrypoint)
        try:
            yield Recording(stream=receive_stream, summary=summary, rate_hz=rate_hz)
        finally:
            # Cancel + drain before the CM returns — producer lifetime
            # is strictly nested inside the ``async with``.
            tg.cancel()

    finished_at = datetime.now(UTC)
    p50, p99 = _tick_percentiles(tick_durations_ms)
    summary.finished_at = finished_at
    summary.tick_duration_ms_p50 = p50
    summary.tick_duration_ms_p99 = p99
    _logger.info(
        "recorder.stop emitted=%s late=%s max_drift_ms=%.3f "
        "tick_p50_ms=%.3f tick_p99_ms=%.3f duration_s=%.3f",
        summary.samples_emitted,
        summary.samples_late,
        summary.max_drift_ms,
        summary.tick_duration_ms_p50,
        summary.tick_duration_ms_p99,
        (finished_at - started_at).total_seconds(),
    )

@asynccontextmanager
async def record(
    source: PollSource,
    *,
    parameters: Sequence[str | int],
    rate_hz: float,
    duration: float | None = None,
    names: Sequence[str] | None = None,
    instances: Sequence[int] = (1,),
    overflow: OverflowPolicy = OverflowPolicy.BLOCK,
    buffer_size: int = 64,
    auto_reconnect: bool = False,
    reconnect_factory: Callable[[], Awaitable[PollSource]] | None = None,
) -> AsyncGenerator[Recording[Sequence[Sample]]]:
    """Record polled samples into a receive stream at an absolute cadence.

    Usage::

        async with record(
            controller, parameters=["process_value", "setpoint"], rate_hz=2, duration=10
        ) as recording:
            async for batch in recording.stream:
                for sample in batch:
                    print(sample.parameter, sample.value)
            # recording.summary is live; recording.summary.finished_at is None
            # while running and set on CM exit.

    The CM yields a :class:`Recording[Sequence[Sample]]` exposing
    ``.stream`` (async iterator of per-tick :class:`Sample` batches),
    ``.summary`` (live :class:`AcquisitionSummary` — recorder is sole
    writer), and ``.rate_hz`` (the cadence the recorder is running
    at).

    Each batch is a flat :class:`Sequence` — one entry per (device,
    parameter, instance) read that succeeded. Failed reads are dropped
    by the source and logged at WARN.

    Args:
        source: Any :class:`PollSource` (a :class:`Controller` or a
            :class:`WatlowManager`).
        parameters: Parameter names or registry IDs to poll each tick.
        rate_hz: Target cadence. Absolute targets are computed
            ``target[n] = start + n * (1 / rate_hz)``. Must be > 0.
        duration: Total acquisition duration in seconds. ``None``
            means "until the caller exits the CM".
        names: Subset of device names to poll per tick. ``None`` polls
            everything the source manages. Ignored for solo controllers.
        instances: 1-indexed loop / channel numbers per device. Single-
            loop devices use ``(1,)``.
        overflow: Backpressure policy when the receive-stream buffer
            is full. See :class:`OverflowPolicy`.
        buffer_size: Receive-stream capacity, in per-tick batches.
        auto_reconnect: When ``True``, treat
            :class:`WatlowConnectionError` raised by ``source.poll_many``
            as a transient transport drop rather than a fatal error.
            The producer logs ``recorder.disconnected``, waits per the
            backoff schedule, and either rebuilds the source via
            ``reconnect_factory`` (if supplied) or simply retries the
            same ``source.poll_many`` on the next tick. ``samples_late``
            ticks up for each tick missed during the gap.
        reconnect_factory: When supplied alongside ``auto_reconnect``,
            invoked to rebuild the :class:`PollSource` after a
            disconnect. Useful when the source's transport needs to be
            re-opened explicitly (e.g. a fresh
            :func:`watlowlib.open_device` call). The returned source
            replaces ``source`` for subsequent ticks. Without a
            factory, the recorder relies on ``source.poll_many`` itself to
            recover (which works for callers that wrap their own
            transport-reopen logic inside ``poll_many``).

    Yields:
        A :class:`Recording[Sequence[Sample]]` exposing ``.stream``,
        ``.summary``, and ``.rate_hz``.

    Raises:
        ValueError: ``rate_hz <= 0``, ``duration <= 0``, or
            ``buffer_size < 1``.
    """
    if rate_hz <= 0:
        raise ValueError(f"rate_hz must be > 0, got {rate_hz!r}")
    if duration is not None and duration <= 0:
        raise ValueError(f"duration must be > 0 or None, got {duration!r}")
    if buffer_size < 1:
        raise ValueError(f"buffer_size must be >= 1, got {buffer_size!r}")
    if not parameters:
        raise ValueError("parameters must be a non-empty sequence")

    period = 1.0 / rate_hz
    total_ticks = None if duration is None else max(1, round(duration * rate_hz))

    send_stream, receive_stream = anyio.create_memory_object_stream[Sequence[Sample]](
        max_buffer_size=buffer_size,
    )
    # Producer-side clone of the receive stream — used to evict the
    # oldest queued batch under DROP_OLDEST. Cloning here (before the
    # consumer starts iterating) keeps the eviction path off the
    # consumer's iterator and avoids racing with it.
    drop_rx = receive_stream.clone()

    started_at = datetime.now(UTC)
    summary = AcquisitionSummary(started_at=started_at)
    tick_durations_ms: list[float] = []
    _logger.info(
        "recorder.start rate_hz=%s duration_s=%s overflow=%s buffer_size=%s names=%s",
        rate_hz,
        duration,
        overflow.value,
        buffer_size,
        list(names) if names is not None else None,
    )

    async with anyio.create_task_group() as tg, receive_stream:

        async def _producer_entrypoint() -> None:
            await _run_producer(
                source,
                send_stream,
                drop_rx,
                tuple(parameters),
                tuple(instances),
                names,
                period,
                total_ticks,
                overflow,
                summary,
                tick_durations_ms,
                auto_reconnect=auto_reconnect,
                reconnect_factory=reconnect_factory,
            )

        # ``start_soon`` returns a TaskHandle (anyio >= 4.14); the producer is
        # fire-and-forget and cancelled via the group below, so discard it.
        _ = tg.start_soon(_producer_entrypoint)
        try:
            yield Recording(stream=receive_stream, summary=summary, rate_hz=rate_hz)
        finally:
            # Cancel + drain before the CM returns — producer lifetime
            # is strictly nested inside the ``async with``.
            tg.cancel()

    finished_at = datetime.now(UTC)
    p50, p99 = _tick_percentiles(tick_durations_ms)
    summary.finished_at = finished_at
    summary.tick_duration_ms_p50 = p50
    summary.tick_duration_ms_p99 = p99
    _logger.info(
        "recorder.stop emitted=%s late=%s max_drift_ms=%.3f "
        "tick_p50_ms=%.3f tick_p99_ms=%.3f duration_s=%.3f",
        summary.samples_emitted,
        summary.samples_late,
        summary.max_drift_ms,
        summary.tick_duration_ms_p50,
        summary.tick_duration_ms_p99,
        (finished_at - started_at).total_seconds(),
    )