"""Shared ADIOS stream helpers."""
from __future__ import annotations
from contextlib import contextmanager
import adios2.stream
from adios2 import Adios, Stream
[docs]
def iter_stream_steps(stream, step_timeout_secs=None):
"""
Iterate ADIOS stream steps with optional timeout-as-EOF handling.
Parameters
----------
stream
ADIOS stream object exposing either normal iteration or
``steps(timeout=...)``.
step_timeout_secs : int, float, or None, optional
Timeout passed to ``stream.steps``. When the timeout expires, ADIOS
raises a runtime error; this helper treats the expected "no new step"
message as end-of-stream for offline files.
Yields
------
object
ADIOS step objects from the stream.
"""
try:
if step_timeout_secs is None:
for step in stream:
yield step
else:
for step in stream.steps(timeout=float(step_timeout_secs)):
yield step
except RuntimeError as exc:
msg = str(exc).lower()
if "no new step within the time limit" in msg:
return
raise
[docs]
@contextmanager
def open_adios_stream(filename, open_timeout_secs=None, *, force_file_stream=True):
"""
Open an ADIOS stream with optional ``OpenTimeoutSecs`` configuration.
Parameters
----------
filename : str
ADIOS-readable source path.
open_timeout_secs : int, float, or None, optional
Open timeout for ``FileStream`` mode. If omitted, the default ADIOS
opening path is used.
force_file_stream : bool, optional
If True and ``open_timeout_secs`` is provided, open with a declared IO
using ``FileStream``. Campaign ``.aca`` sources should normally leave
this False or use ``adios2.FileReader`` instead.
Yields
------
adios2.Stream
Open stream context.
"""
if open_timeout_secs is None or not force_file_stream:
with adios2.stream.Stream(filename, "r") as stream:
yield stream
return
adios = Adios()
io = adios.declare_io(f"read_bp_{id(filename)}")
io.set_engine("FileStream")
io.set_parameter("OpenTimeoutSecs", str(open_timeout_secs))
with Stream(io, filename, "r") as stream:
yield stream