aimbat.io

def clear_seismogram_data_cache() -> None:
    """Drop every entry from the in-memory waveform read cache.

    Staged (uncommitted) writes in `_pending` are left intact.
    """
    _cache.clear()

def create_event(datasource: str | PathLike[str], datatype: DataType) -> AimbatEvent:
    """Create an `AimbatEvent` from a data source.

    Args:
        datasource: Logical source identifier for the data source.
        datatype: Data type of the source.

    Returns:
        A new `AimbatEvent` instance.

    Raises:
        NotImplementedError: If `datatype` has no registered event creator.
    """
    logger.debug(f"Creating AimbatEvent from {datasource}.")
    creator = _event_creators.get(datatype)
    if creator is None:
        raise NotImplementedError(f"{datatype} does not support event creation.")
    return creator(datasource)

def create_seismogram(
    datasource: str | PathLike[str], datatype: DataType
) -> AimbatSeismogram:
    """Create an `AimbatSeismogram` from a data source.

    Args:
        datasource: Logical source identifier for the data source.
        datatype: Data type of the source.

    Returns:
        A new `AimbatSeismogram` instance.

    Raises:
        NotImplementedError: If `datatype` has no registered seismogram creator.
    """
    logger.debug(f"Creating AimbatSeismogram from {datasource}.")
    creator = _seismogram_creators.get(datatype)
    if creator is None:
        raise NotImplementedError(f"{datatype} does not support seismogram creation.")
    return creator(datasource)

def create_station(
    datasource: str | PathLike[str], datatype: DataType
) -> AimbatStation:
    """Create an `AimbatStation` from a data source.

    Args:
        datasource: Logical source identifier for the data source.
        datatype: Data type of the source.

    Returns:
        A new `AimbatStation` instance.

    Raises:
        NotImplementedError: If `datatype` has no registered station creator.
    """
    logger.debug(f"Creating AimbatStation from {datasource}.")
    creator = _station_creators.get(datatype)
    if creator is None:
        raise NotImplementedError(f"{datatype} does not support station creation.")
    return creator(datasource)

def event_creator(
    datatype: DataType,
) -> Callable[
    [Callable[[str | PathLike[str]], AimbatEvent]],
    Callable[[str | PathLike[str]], AimbatEvent],
]:
    """Decorator that registers a function as an event creator for `datatype`.

    Args:
        datatype: The data type the decorated function creates events from.

    Example:
        ```python
        @event_creator(DataType.SAC)
        def create_event_from_sacfile(sacfile: str | PathLike[str]) -> AimbatEvent: ...
        ```
    """

    def decorator(
        fn: Callable[[str | PathLike[str]], AimbatEvent],
    ) -> Callable[[str | PathLike[str]], AimbatEvent]:
        register_event_creator(datatype, fn)
        return fn

    return decorator

def file_source_present(sourcename: str | PathLike[str]) -> bool:
    """Return whether a filesystem-backed data source is present.

    Shared implementation for the file-based probes (SAC, miniSEED, JSON). A
    reachable regular file is present; a reachable directory that simply does
    not contain the named entry is a clean `False`. Anything that leaves the
    answer unknowable right now - a missing or unreadable parent directory (a
    moved or unmounted tree), a permission error, an unreachable mount -
    raises `SourceUnavailableError` rather than being mistaken for a deletion.

    Args:
        sourcename: Path to the file backing the data source.

    Raises:
        SourceUnavailableError: If presence cannot be determined right now.
    """
    path = Path(sourcename)
    try:
        mode = os.stat(path).st_mode
    except FileNotFoundError as exc:
        # The file is gone. Only trust that as a real deletion if the parent
        # directory is still reachable; an unmounted volume whose mountpoint
        # lingers would otherwise look like a clean absence.
        if os.path.isdir(path.parent):
            return False
        raise SourceUnavailableError(
            f"Neither {os.fspath(path)!r} nor its parent directory could be "
            + "reached; this looks like a moved or unmounted data tree, not a "
            + "deletion."
        ) from exc
    except OSError as exc:
        raise SourceUnavailableError(
            f"Could not determine whether {os.fspath(path)!r} is present: {exc}"
        ) from exc
    return stat_module.S_ISREG(mode)

def read_seismogram_data(
    datasource: str | PathLike[str],
    datatype: DataType,
    session: Session | None = None,
) -> npt.NDArray[np.floating]:
    """Read seismogram waveform data from a data source.

    Results are cached in memory by `(datasource, datatype)` key, alongside
    the source's modification time and size; a cached array is re-read when
    either has changed since, so a source written behind AIMBAT's back does
    not keep serving stale data. The returned array is read-only.

    If `session` is given and has a staged write for this key (from
    `stage_seismogram_data`), the staged value is returned instead of reading
    the data source, so a session sees its own uncommitted waveform write.

    Args:
        datasource: Logical source identifier for the data source.
        datatype: Data type of the source.
        session: Session whose staged writes should be consulted first.

    Returns:
        Read-only seismogram waveform data as a NumPy array.

    Raises:
        NotImplementedError: If `datatype` has no registered data reader.
    """
    logger.debug(f"Reading seismogram data from {datasource}.")
    key = (str(datasource), datatype)
    if session is not None:
        # Checked before the reader-registration check below: a session's own
        # staged write should be readable back regardless of whether the
        # datatype also supports reading through a registered reader.
        staged = _pending.get(session)
        if staged is not None and key in staged:
            logger.debug(f"Retrieved staged seismogram data for {datasource}.")
            return staged[key]
    reader = _seismogram_data_readers.get(datatype)
    if reader is None:
        raise NotImplementedError(
            f"{datatype} does not support reading seismogram data."
        )
    stamp = _source_stamp(key[0])
    entry = _cache.get(key)
    if entry is not None and entry.stamp == stamp:
        logger.debug(f"Retrieved seismogram data from cache for {datasource}.")
        _cache.move_to_end(key)
        return entry.data

    context = SeismogramReadContext(
        sourcename=key[0], datatype=datatype, session=session
    )
    arr = reader(context)
    arr.flags.writeable = False
    _cache[key] = _CacheEntry(stamp, arr)
    _cache.move_to_end(key)
    if len(_cache) > _CACHE_MAX_ENTRIES:
        _cache.popitem(last=False)
    return arr

def read_source_records(
    datasource: str | PathLike[str],
    datatype: DataType,
    *,
    station: bool = True,
    event: bool = True,
    seismogram: bool = True,
) -> SourceRecords:
    """Build the requested records from a data source, reading it once where possible.

    Data types with a registered source-records reader (SAC) read the source
    a single time and build everything asked for from it. The rest fall back
    to their individual creators, one read each.

    Args:
        datasource: Logical source identifier for the data source.
        datatype: Data type of the source.
        station: Whether an `AimbatStation` is wanted.
        event: Whether an `AimbatEvent` is wanted.
        seismogram: Whether an `AimbatSeismogram` is wanted.

    Returns:
        The requested records; the fields that were not asked for are `None`.

    Raises:
        NotImplementedError: If a record is asked for that `datatype` has no
            registered creator for.
    """
    if not (station or event or seismogram):
        return SourceRecords()

    reader = _source_records_readers.get(datatype)
    if reader is not None:
        return reader(
            SourceRecordsRequest(
                sourcename=str(datasource),
                datatype=datatype,
                station=station,
                event=event,
                seismogram=seismogram,
            )
        )

    return SourceRecords(
        station=create_station(datasource, datatype) if station else None,
        event=create_event(datasource, datatype) if event else None,
        seismogram=create_seismogram(datasource, datatype) if seismogram else None,
    )

def register_event_creator(
    datatype: DataType,
    fn: Callable[[str | PathLike[str]], AimbatEvent],
) -> None:
    """Register a function that creates an `AimbatEvent` from a data source.

    Args:
        datatype: The data type this creator handles.
        fn: Callable that accepts a logical source identifier and returns an
            `AimbatEvent` instance.
    """
    if datatype in _event_creators:
        logger.warning(f"Overwriting existing event creator for {datatype}.")
    logger.debug(f"Registering event creator for {datatype}.")
    _event_creators[datatype] = fn

def register_seismogram_creator(
    datatype: DataType,
    fn: Callable[[str | PathLike[str]], AimbatSeismogram],
) -> None:
    """Register a function that creates an `AimbatSeismogram` from a data source.

    Args:
        datatype: The data type this creator handles.
        fn: Callable that accepts a logical source identifier and returns an
            `AimbatSeismogram` instance.
    """
    if datatype in _seismogram_creators:
        logger.warning(f"Overwriting existing seismogram creator for {datatype}.")
    logger.debug(f"Registering seismogram creator for {datatype}.")
    _seismogram_creators[datatype] = fn

def register_seismogram_data_reader(
    datatype: DataType,
    fn: SeismogramDataReader,
) -> None:
    """Register a function that reads seismogram waveform data from a data source.

    Args:
        datatype: The data type this reader handles.
        fn: Callable that accepts a `SeismogramReadContext` and returns the
            waveform data as a NumPy array.
    """
    if datatype in _seismogram_data_readers:
        logger.warning(f"Overwriting existing seismogram data reader for {datatype}.")
    logger.debug(f"Registering seismogram data reader for {datatype}.")
    _seismogram_data_readers[datatype] = fn

def register_seismogram_data_writer(
    datatype: DataType,
    fn: Callable[[str | PathLike[str], npt.NDArray[np.floating]], None],
) -> None:
    """Register a function that writes seismogram waveform data to a data source.

    Args:
        datatype: The data type this writer handles.
        fn: Callable that accepts a logical source identifier and a NumPy array,
            and writes the data to the source.
    """
    if datatype in _seismogram_data_writers:
        logger.warning(f"Overwriting existing seismogram data writer for {datatype}.")
    logger.debug(f"Registering seismogram data writer for {datatype}.")
    _seismogram_data_writers[datatype] = fn

def register_source_probe(
    datatype: DataType,
    fn: SourcePresenceProbe,
) -> None:
    """Register a function that checks whether a data source is still present.

    Args:
        datatype: The data type this probe handles.
        fn: Callable that accepts a `SeismogramReadContext` and returns whether
            the source still provides the seismogram, raising
            `SourceUnavailableError` when it cannot tell.
    """
    if datatype in _source_probes:
        logger.warning(f"Overwriting existing source presence probe for {datatype}.")
    logger.debug(f"Registering source presence probe for {datatype}.")
    _source_probes[datatype] = fn

def register_source_records_reader(
    datatype: DataType,
    fn: SourceRecordsReader,
) -> None:
    """Register a function that builds several records from one read of a data source.

    Optional: a data type with no reader of its own is served by calling its
    individual creators instead (see `read_source_records`). Registering one
    is worth it when a single read yields more than one record, as it does
    for SAC.

    Args:
        datatype: The data type this reader handles.
        fn: Callable that accepts a `SourceRecordsRequest` and returns the
            requested records as a `SourceRecords`.
    """
    if datatype in _source_records_readers:
        logger.warning(f"Overwriting existing source records reader for {datatype}.")
    logger.debug(f"Registering source records reader for {datatype}.")
    _source_records_readers[datatype] = fn

def register_station_creator(
    datatype: DataType,
    fn: Callable[[str | PathLike[str]], AimbatStation],
) -> None:
    """Register a function that creates an `AimbatStation` from a data source.

    Args:
        datatype: The data type this creator handles.
        fn: Callable that accepts a logical source identifier and returns an
            `AimbatStation` instance.
    """
    if datatype in _station_creators:
        logger.warning(f"Overwriting existing station creator for {datatype}.")
    logger.debug(f"Registering station creator for {datatype}.")
    _station_creators[datatype] = fn

def seismogram_creator(
    datatype: DataType,
) -> Callable[
    [Callable[[str | PathLike[str]], AimbatSeismogram]],
    Callable[[str | PathLike[str]], AimbatSeismogram],
]:
    """Decorator that registers a function as a seismogram creator for `datatype`.

    Args:
        datatype: The data type the decorated function creates seismograms from.

    Example:
        ```python
        @seismogram_creator(DataType.SAC)
        def create_seismogram_from_sacfile(
            sacfile: str | PathLike[str],
        ) -> AimbatSeismogram: ...
        ```
    """

    def decorator(
        fn: Callable[[str | PathLike[str]], AimbatSeismogram],
    ) -> Callable[[str | PathLike[str]], AimbatSeismogram]:
        register_seismogram_creator(datatype, fn)
        return fn

    return decorator

def seismogram_data_reader(
    datatype: DataType,
) -> Callable[[SeismogramDataReader], SeismogramDataReader]:
    """Decorator that registers a function as a seismogram data reader for `datatype`.

    Args:
        datatype: The data type the decorated function reads waveform data from.

    Example:
        ```python
        @seismogram_data_reader(DataType.SAC)
        def read_seismogram_data_from_sacfile(
            context: SeismogramReadContext,
        ) -> npt.NDArray[np.floating]: ...
        ```
    """

    def decorator(fn: SeismogramDataReader) -> SeismogramDataReader:
        register_seismogram_data_reader(datatype, fn)
        return fn

    return decorator

def seismogram_data_writer(
    datatype: DataType,
) -> Callable[
    [Callable[[str | PathLike[str], npt.NDArray[np.floating]], None]],
    Callable[[str | PathLike[str], npt.NDArray[np.floating]], None],
]:
    """Decorator that registers a function as a seismogram data writer for `datatype`.

    Args:
        datatype: The data type the decorated function writes waveform data to.

    Example:
        ```python
        @seismogram_data_writer(DataType.SAC)
        def write_seismogram_data_to_sacfile(
            sacfile: str | PathLike[str], data: npt.NDArray[np.floating]
        ) -> None: ...
        ```
    """

    def decorator(
        fn: Callable[[str | PathLike[str], npt.NDArray[np.floating]], None],
    ) -> Callable[[str | PathLike[str], npt.NDArray[np.floating]], None]:
        register_seismogram_data_writer(datatype, fn)
        return fn

    return decorator

def source_present(
    sourcename: str, datatype: DataType, session: Session | None = None
) -> bool:
    """Check whether a data source is still present and provides its seismogram.

    Args:
        sourcename: Logical source identifier for the data source.
        datatype: Data type of the source.
        session: Session made available to the probe (needed by sources whose
            identity is resolved through the database).

    Returns:
        `True` if the source is reachable and provides the seismogram,
        `False` if it is reachable and definitely does not.

    Raises:
        NotImplementedError: If `datatype` has no registered probe.
        SourceUnavailableError: If the probe cannot determine presence right
            now (e.g. an unreadable file or an unreachable mount).
    """
    probe = _source_probes.get(datatype)
    if probe is None:
        raise NotImplementedError(f"{datatype} does not support source probing.")
    context = SeismogramReadContext(
        sourcename=str(sourcename), datatype=datatype, session=session
    )
    return probe(context)

def source_probe(
    datatype: DataType,
) -> Callable[[SourcePresenceProbe], SourcePresenceProbe]:
    """Decorator that registers a function as a source-presence probe for `datatype`.

    Args:
        datatype: The data type the decorated function probes.

    Example:
        ```python
        @source_probe(DataType.SAC)
        def sac_source_present(context: SeismogramReadContext) -> bool: ...
        ```
    """

    def decorator(fn: SourcePresenceProbe) -> SourcePresenceProbe:
        register_source_probe(datatype, fn)
        return fn

    return decorator

def source_records_reader(
    datatype: DataType,
) -> Callable[[SourceRecordsReader], SourceRecordsReader]:
    """Decorator that registers a function as a source-records reader for `datatype`.

    Args:
        datatype: The data type the decorated function reads records from.

    Example:
        ```python
        @source_records_reader(DataType.SAC)
        def read_records_from_sacfile(
            request: SourceRecordsRequest,
        ) -> SourceRecords: ...
        ```
    """

    def decorator(fn: SourceRecordsReader) -> SourceRecordsReader:
        register_source_records_reader(datatype, fn)
        return fn

    return decorator

def stage_seismogram_data(
    session: Session | None,
    datasource: str | PathLike[str],
    datatype: DataType,
    data: npt.NDArray[np.floating],
) -> None:
    """Stage a waveform write to flush when `session` commits.

    The write is validated immediately (a datatype with no registered writer
    raises straight away), but the data source is not touched until `session`
    commits. The staged value is discarded if `session` rolls back or is
    abandoned, and is visible to reads made through the same session before
    the commit.

    A copy of `data` is taken at call time, so later mutation of the caller's
    array does not change what is persisted. The `(datasource, datatype)` key
    is also captured now: renaming the data source on the model before the
    commit would flush to the old path. Staged pages are held in memory until
    the commit with no eviction (unlike the read cache), so staging waveform
    writes for a very large event holds every new array at once.

    Staging is keyed by `(datasource, datatype)`, not by seismogram: staging
    a second write for the same data source in one session before committing
    replaces the first, and only the last survives the flush.

    If `session` is `None` (a detached instance, attached to no session)
    there is no transaction to bind to, so the write falls back to an eager
    `write_seismogram_data` call.

    Args:
        session: The session that owns the seismogram, or `None`.
        datasource: Logical source identifier for the data source.
        datatype: Data type of the source.
        data: Seismogram waveform data to write.

    Raises:
        NotImplementedError: If `datatype` has no registered data writer.
    """
    if datatype not in _seismogram_data_writers:
        raise NotImplementedError(
            f"{datatype} does not support writing seismogram data."
        )
    if session is None:
        msg = f"Waveform write for detached seismogram {datasource}: "
        logger.warning(msg + "no session to bind to, writing eagerly.")
        write_seismogram_data(datasource, datatype, data)
        return
    logger.debug(f"Staging seismogram data write to {datasource}.")
    staged = np.array(data, copy=True)
    staged.flags.writeable = False
    session_pending = _pending.setdefault(session, {})
    key = (str(datasource), datatype)
    if key in session_pending:
        logger.warning(
            f"Replacing an earlier staged waveform write for {datasource}; "
            + "only the latest is flushed on commit."
        )
    session_pending[key] = staged