fujilib.sinks¶

async def pipe(
    source: Recording[Batch] | MemoryObjectReceiveStream[Batch],
    sink: SampleSink,
    *,
    batch_size: int = 64,
    flush_interval: float = 1.0,
) -> AcquisitionSummary:
    """Write every batch of a recording to the open ``sink``, until the stream ends.

    Samples are written in groups: once ``batch_size`` are waiting, and at the
    latest ``flush_interval`` seconds after the first of them arrived, even if
    no more come. When the pipe stops, by the stream ending, an error or
    cancellation, it takes the batches already waiting in the stream and writes
    everything it holds first (for up to 30 s, even when cancelled), unless the
    sink itself failed. So a recording stopped with Ctrl-C has a row for every
    poll its summary counts.

    Returns:
        What the pipe wrote: ``samples_emitted`` counts batches (polls), and
        ``error_samples`` the samples of failed polls. The recording's own
        counters are in its ``summary``.

    Raises:
        FujiValidationError: ``source`` is not a recording or a memory object
            stream, or ``batch_size`` or ``flush_interval`` is invalid.
        FujiSinkError: the sink failed.
        FujiSinkWriteError: the last samples could not be written in time.
    """
    stream = source.stream if isinstance(source, Recording) else source
    if not _is_memory_stream(stream):
        msg = f"source must be a Recording or its stream, got {type(source).__name__}"
        raise FujiValidationError(msg)
    if not _is_count(batch_size) or batch_size < 1:
        msg = f"batch_size must be an integer of at least 1, got {batch_size!r}"
        raise FujiValidationError(msg)
    if not _is_seconds(flush_interval) or flush_interval <= 0:
        msg = f"flush_interval must be a finite number of seconds above 0, got {flush_interval!r}"
        raise FujiValidationError(msg)
    writer = _PipeWriter(sink)
    try:
        await writer.run(stream, batch_size, flush_interval)
    finally:
        writer.drain(stream)
        await writer.finish()
    if writer.unwritten:
        msg = f"could not write the last {writer.unwritten} samples within {_FINAL_WRITE_S:g} s"
        raise FujiSinkWriteError(msg)
    return writer.summary

def row_columns(channels: Iterable[ChannelId]) -> tuple[ColumnSpec, ...]:
    """The columns of every row for an analyzer with these established ``channels``.

    Reading and analyzer columns are nullable, because an error row carries
    ``None`` in all of them.
    """
    specs = list(HEADER_COLUMNS)
    for channel in channels:
        prefix = _prefix(channel)
        specs.extend(
            ColumnSpec(prefix + c.name, c.python_type, nullable=True) for c in READING_COLUMNS
        )
    specs.extend(ColumnSpec(c.name, c.python_type, nullable=True) for c in ANALYZER_COLUMNS)
    specs.extend(_ERROR_COLUMNS)
    return tuple(specs)

def sample_channels(sample: Sample) -> tuple[ChannelId, ...]:
    """The channels ``sample``'s row has columns for.

    :attr:`Sample.channels <fujilib.streaming.sample.Sample.channels>`, or,
    for a sample made without them, its frame's.
    """
    if sample.channels:
        return sample.channels
    return sample.frame.channels if sample.frame is not None else ()

def sample_to_row(
    sample: Sample,
    channels: Iterable[ChannelId] | None = None,
) -> dict[str, Scalar]:
    """Flatten ``sample`` into one wide row; see the module docstring for the layout.

    Args:
        sample: The sample to flatten.
        channels: The channels that fix the row's columns. When omitted, the
            sample's own (:func:`sample_channels`): a recorder sets them on
            every sample, failed polls included, so every row of a recording
            has the same keys.
    """
    frame = sample.frame
    established = tuple(channels) if channels is not None else sample_channels(sample)
    row: dict[str, Scalar] = {
        "device": sample.device,
        "address": sample.address,
        "protocol": sample.protocol.value,
        "t_mono_ns": sample.t_mono_ns,
        "t_utc": sample.t_utc.isoformat(),
        "t_midpoint_mono_ns": sample.t_midpoint_mono_ns,
        "requested_at": sample.requested_at.isoformat(),
        "received_at": sample.received_at.isoformat(),
        "latency_s": sample.latency_s,
    }
    readings = {r.channel: r for r in frame.readings} if frame is not None else {}
    for channel in established:
        prefix = _prefix(channel)
        reading = readings.get(channel)
        values = reading.as_dict() if reading is not None else {}
        row.update({prefix + c.name: values.get(c.name) for c in READING_COLUMNS})
    analyzer = frame.analyzer.as_dict() if frame is not None and frame.analyzer is not None else {}
    row.update({c.name: analyzer.get(c.name) for c in ANALYZER_COLUMNS})
    error = sample.error
    if error is not None:
        cls = type(error)
        row["error_type"] = f"{cls.__module__}.{cls.__qualname__}"
        row["error_message"] = str(error)
    else:
        row["error_type"] = None
        row["error_message"] = None
    return row

def csv_cell(value: Scalar) -> str | int | float | None:
    """A row value as the CSV writer takes it: a boolean becomes ``true`` / ``false``."""
    if isinstance(value, bool):
        return "true" if value else "false"
    return value

def require_pyarrow() -> None:
    """Check that ``pyarrow`` can be imported.

    Raises:
        FujiSinkDependencyError: it cannot; the ``parquet`` extra is missing.
    """
    try:
        _ = importlib.import_module("pyarrow")
        _ = importlib.import_module("pyarrow.parquet")
    except ImportError as exc:
        msg = f"the Parquet sink needs pyarrow, which is not installed; {_INSTALL_HINT}"
        raise FujiSinkDependencyError(msg) from exc