Source code for signal_dataset.storage.staging.passthrough

"""Staging for locations the native reader can already open.

Local paths and ``gs://`` URIs are handed to ArrayRecord unchanged. This also
owns the fragment convention that pins a GCS generation, so that the literal
lives in exactly one place rather than at every reader construction.
"""

from __future__ import annotations

from collections.abc import Iterator
from contextlib import contextmanager

from signal_dataset.storage.staging.contracts import ShardStagingArea




class PassthroughStaging(ShardStagingArea):
    """Yields the URI itself, pinning a generation when one is given."""



    @contextmanager
    def stage(self, uri: str, *, generation: int | None = None) -> Iterator[str]:
        yield uri if generation is None else f"{uri}#{generation}"