rustwx_stormscope/datasources.py

#!/usr/bin/env python
"""GOES/MRMS data backends for StormScope.

Two backends:
  - 'earth2studio' (default): uses earth2studio's built-in GOES/MRMS data sources (Python S3 fetch).
  - 'rust': uses the rustwx-backed `stormscope_obs` extension for parallel download + pure-Rust decode,
    exposed through earth2studio-compatible RustGOES / RustMRMS adapters (see M2).

`install_backend(core, backend)` swaps the data-source classes that `core.run_stormscope` uses, so the
rest of the pipeline (interpolation, sampling, output) is unchanged.
"""


def install_backend(core, backend: str):
    """Select the GOES/MRMS data backend on `core`. Falls back to earth2studio if rust is missing."""
    if backend == "earth2studio":
        core.RUST_BACKEND = False
        return
    if backend == "rust":
        try:
            import stormscope_obs  # noqa: F401  (compiled rust extension)
            import h5netcdf  # noqa: F401      (GOES decode)
            core.RUST_BACKEND = True
            print("[rustwx_stormscope] rust data backend active (parallel download + MRMS rust decode)", flush=True)
        except Exception as e:
            core.RUST_BACKEND = False
            print(f"[rustwx_stormscope] rust backend unavailable ({type(e).__name__}); "
                  f"falling back to earth2studio fetch", flush=True)
        return
    raise ValueError(f"unknown backend {backend!r}")