qpkit 0.1.0__py3-none-any.whl

qpkit/__init__.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from .api import QPKit, Request
+from .exceptions import DownloadError, PlanBuildError, QPKitError
+from .http import HttpClient, HttpClientProtocol
+from .models import (
+    BoundingBox,
+    DownloadDSSResult,
+    DownloadItem,
+    DownloadPlan,
+    DownloadResult,
+    DSSWriteResult,
+    GridOptions,
+    HRRRGridOptions,
+    HRRRRequest,
+    QPEGridOptions,
+    QPERequest,
+    QPFGridOptions,
+    QPFRequest,
+)
+
+try:
+    from ._version import version as __version__
+except ImportError:
+    try:
+        from importlib.metadata import version
+
+        __version__ = version("qpkit")
+    except Exception:
+        __version__ = "0.0.0"
+
+__all__ = [
+    "BoundingBox",
+    "DSSWriteResult",
+    "DownloadDSSResult",
+    "DownloadError",
+    "DownloadItem",
+    "DownloadPlan",
+    "DownloadResult",
+    "GridOptions",
+    "HRRRGridOptions",
+    "HRRRRequest",
+    "HttpClient",
+    "QPEGridOptions",
+    "HttpClientProtocol",
+    "PlanBuildError",
+    "QPERequest",
+    "QPFGridOptions",
+    "QPFRequest",
+    "QPKit",
+    "QPKitError",
+    "Request",
+    "__version__",
+]

qpkit/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

qpkit/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

qpkit/api.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+import logging
+import os
+import warnings
+from datetime import datetime, timezone
+from pathlib import Path
+from types import TracebackType
+
+from .downloader import Downloader
+from .http import HttpClient, HttpClientProtocol
+from .models import (
+    DownloadDSSResult,
+    DownloadPlan,
+    DownloadResult,
+    GridOptions,
+    HRRRGridOptions,
+    HRRRRequest,
+    QPEGridOptions,
+    QPERequest,
+    QPFGridOptions,
+    QPFRequest,
+)
+from .services import (
+    HRRRDSSWriter,
+    HRRRService,
+    QPEDSSWriter,
+    QPEService,
+    QPFDSSWriter,
+    QPFService,
+)
+from .utils import list_local_files
+
+Request = QPERequest | QPFRequest | HRRRRequest
+
+
+class QPKit:
+    """High-level API for building plans and downloading precipitation products."""
+
+    def __init__(
+        self,
+        *,
+        http_client: HttpClientProtocol | None = None,
+        logger: logging.Logger | None = None,
+        max_workers: int = 4,
+    ) -> None:
+        self._owns_http = http_client is None
+        self.http: HttpClientProtocol = http_client or HttpClient()
+        self.logger = logger or logging.getLogger("qpkit")
+        self.qpe = QPEService(self.http)
+        self.qpf = QPFService(self.http)
+        self.hrrr = HRRRService(self.http)
+        self.downloader = Downloader(self.http, self.logger, max_workers=max_workers)
+        self.hrrr_dss = HRRRDSSWriter()
+        self.qpe_dss = QPEDSSWriter()
+        self.qpf_dss = QPFDSSWriter()
+
+    def close(self) -> None:
+        if self._owns_http:
+            self.http.close()
+
+    def __enter__(self) -> QPKit:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    @staticmethod
+    def _resolve(output_dir: str | os.PathLike[str]) -> Path:
+        return Path(os.fspath(output_dir))
+
+    def build_plan_qpe(
+        self,
+        request: QPERequest,
+        output_dir: str | os.PathLike[str],
+        *,
+        force: bool = False,
+    ) -> DownloadPlan:
+        path = self._resolve(output_dir)
+        return self.qpe.build_plan(request, path, list_local_files(path), force=force)
+
+    def build_plan_qpf(
+        self,
+        request: QPFRequest,
+        output_dir: str | os.PathLike[str],
+        *,
+        force: bool = False,
+    ) -> DownloadPlan:
+        path = self._resolve(output_dir)
+        return self.qpf.build_plan(request, path, list_local_files(path), force=force)
+
+    def build_plan_hrrr(
+        self,
+        request: HRRRRequest,
+        output_dir: str | os.PathLike[str],
+        *,
+        force: bool = False,
+    ) -> DownloadPlan:
+        path = self._resolve(output_dir)
+        return self.hrrr.build_plan(request, path, list_local_files(path), force=force)
+
+    def download_qpe(
+        self,
+        request: QPERequest,
+        output_dir: str | os.PathLike[str],
+        *,
+        force: bool = False,
+    ) -> DownloadResult:
+        plan = self.build_plan_qpe(request, output_dir, force=force)
+        return self.downloader.download(plan)
+
+    def download_qpf(
+        self,
+        request: QPFRequest,
+        output_dir: str | os.PathLike[str],
+        *,
+        force: bool = False,
+    ) -> DownloadResult:
+        plan = self.build_plan_qpf(request, output_dir, force=force)
+        return self.downloader.download(plan)
+
+    def download_hrrr(
+        self,
+        request: HRRRRequest,
+        output_dir: str | os.PathLike[str],
+        *,
+        force: bool = False,
+        cycle_hour_latest: bool | None = None,
+    ) -> DownloadResult:
+        if cycle_hour_latest:
+            today = datetime.now(timezone.utc).date()
+            if request.cycle_date is not None and request.cycle_date != today:
+                warnings.warn(
+                    "cycle_hour_latest=True requires cycle_date to be unset or today's "
+                    f"UTC date; got {request.cycle_date}. Using the latest available "
+                    "cycle instead.",
+                    stacklevel=2,
+                )
+            latest_date, latest_hour = self.hrrr.find_latest_cycle()
+            request = request.model_copy(
+                update={"cycle_date": latest_date, "cycle_hour": latest_hour}
+            )
+        plan = self.build_plan_hrrr(request, output_dir, force=force)
+        return self.downloader.download(plan)
+
+    def build_plan(
+        self,
+        request: Request,
+        output_dir: str | os.PathLike[str],
+        *,
+        force: bool = False,
+    ) -> DownloadPlan:
+        match request:
+            case QPERequest():
+                return self.build_plan_qpe(request, output_dir, force=force)
+            case QPFRequest():
+                return self.build_plan_qpf(request, output_dir, force=force)
+            case HRRRRequest():
+                return self.build_plan_hrrr(request, output_dir, force=force)
+
+    def download(
+        self,
+        request: Request,
+        output_dir: str | os.PathLike[str],
+        *,
+        force: bool = False,
+    ) -> DownloadResult:
+        plan = self.build_plan(request, output_dir, force=force)
+        return self.downloader.download(plan)
+
+    def download_to_dss(
+        self,
+        request: Request,
+        output_dir: str | os.PathLike[str],
+        dss_file: str | os.PathLike[str],
+        *,
+        grid_options: GridOptions | None = None,
+        force: bool = False,
+        cycle_hour_latest: bool | None = None,
+        dry_run: bool = False,
+        dss_version: int | None = None,
+        interval: int | None = None,
+    ) -> DownloadDSSResult:
+        """Download a precipitation product and write the resulting grids to DSS.
+
+        Single end-to-end entry point for every supported product. The request
+        type drives the download dispatch; the ``grid_options`` type drives the
+        DSS-writer dispatch. When ``grid_options`` is None, sensible defaults
+        are chosen based on the product.
+
+        ``interval`` (HRRR APCP only) — compute per-interval accumulations
+        instead of storing the raw running total. Provide a single forecast hour
+        in ``request.forecast_hours`` (the end hour); the required intermediate
+        files are derived and downloaded automatically. Not applicable to QPE
+        (its accumulation window is set via ``QPERequest.interval``).
+
+        Two-stage workflows (writing GRIBs that came from elsewhere, or
+        re-writing the same set with different conventions) go through the
+        public writer-service attribute, e.g. ``kit.hrrr_dss.write(...)``.
+        """
+        if dss_version not in (None, 6, 7):
+            raise ValueError(f"dss_version must be 6, 7, or None; got {dss_version!r}")
+
+        # Argument coherence — surface the most specific user mistake first.
+        if cycle_hour_latest and not isinstance(request, HRRRRequest):
+            raise ValueError(
+                "cycle_hour_latest is HRRR-only; "
+                f"got request of type {type(request).__name__}"
+            )
+        if interval is not None and isinstance(request, QPERequest):
+            raise ValueError(
+                "interval is not applicable to QPE; the accumulation window is "
+                "set via QPERequest.interval"
+            )
+        if grid_options is not None and not _grid_options_matches(request, grid_options):
+            raise ValueError(
+                f"grid_options type {type(grid_options).__name__} "
+                f"does not match request type {type(request).__name__}"
+            )
+
+        dss_path = Path(os.fspath(dss_file))
+
+        match request:
+            case HRRRRequest():
+                if interval is not None:
+                    if request.precip_type != "APCP":
+                        raise ValueError(
+                            f"interval is only valid for APCP; "
+                            f"got precip_type={request.precip_type!r}"
+                        )
+                    if len(request.forecast_hours) != 1:
+                        raise ValueError(
+                            "interval mode requires exactly one forecast hour "
+                            f"(the end hour); got {len(request.forecast_hours)}"
+                        )
+                    end_hour = request.forecast_hours[0]
+                    if end_hour % interval != 0:
+                        warnings.warn(
+                            f"end_hour={end_hour} is not divisible by "
+                            f"interval={interval}; the last incomplete interval "
+                            "will be omitted.",
+                            stacklevel=2,
+                        )
+                    download_hours = tuple(range(0, end_hour + 1, interval))
+                    request = request.model_copy(
+                        update={"forecast_hours": download_hours}
+                    )
+
+                if grid_options is None:
+                    if request.precip_type == "APCP":
+                        grid_options = HRRRGridOptions.for_apcp()
+                    else:
+                        grid_options = HRRRGridOptions.for_prate(as_depth=False)
+
+                download = self.download_hrrr(
+                    request, output_dir, force=force,
+                    cycle_hour_latest=cycle_hour_latest,
+                )
+                dss = self.hrrr_dss.write(
+                    download.all_files,
+                    dss_path,
+                    grid_options=grid_options,
+                    dry_run=dry_run,
+                    dss_version=dss_version,
+                    interval=interval,
+                )
+
+            case QPERequest():
+                if grid_options is None:
+                    grid_options = QPEGridOptions()
+
+                download = self.download_qpe(request, output_dir, force=force)
+                dss = self.qpe_dss.write(
+                    download.all_files,
+                    dss_path,
+                    grid_options=grid_options,
+                    dry_run=dry_run,
+                    dss_version=dss_version,
+                )
+
+            case QPFRequest():
+                if grid_options is None:
+                    grid_options = QPFGridOptions()
+
+                download = self.download_qpf(request, output_dir, force=force)
+                dss = self.qpf_dss.write(
+                    download.all_files,
+                    dss_path,
+                    grid_options=grid_options,
+                    dry_run=dry_run,
+                    dss_version=dss_version,
+                )
+
+        return DownloadDSSResult(download=download, dss=dss)
+
+
+def _grid_options_matches(request: Request, grid_options: GridOptions) -> bool:
+    """Return True if the grid_options type is consistent with the request type."""
+    match request:
+        case HRRRRequest():
+            return isinstance(grid_options, HRRRGridOptions)
+        case QPERequest():
+            return isinstance(grid_options, QPEGridOptions)
+        case QPFRequest():
+            return isinstance(grid_options, QPFGridOptions)

qpkit/cli.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import argparse
+import datetime as dt
+import json
+import os
+from collections.abc import Sequence
+from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
+from typing import Any
+
+from pydantic import ValidationError
+
+from .api import QPKit, Request
+from .exceptions import QPKitError
+from .logging_utils import build_logger
+from .models import (
+    BoundingBox,
+    DownloadResult,
+    HRRRRequest,
+    QPERequest,
+    QPFRequest,
+)
+from .utils import parse_forecast_hours
+
+
+def _qpkit_version() -> str:
+    try:
+        return version("qpkit")
+    except PackageNotFoundError:
+        return "0.0.0+unknown"
+
+
+class PathExpandAction(argparse.Action):
+    def __init__(
+        self,
+        option_strings: Sequence[str],
+        dest: str,
+        nargs: int | str | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if nargs is not None:
+            raise ValueError('"nargs" is not supported for PathExpandAction')
+        super().__init__(option_strings, dest, **kwargs)
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: str | Sequence[Any] | None,
+        option_string: str | None = None,
+    ) -> None:
+        if values == "-":
+            setattr(namespace, self.dest, None)
+            return
+        if not isinstance(values, str):
+            raise ValueError("PathExpandAction expects a string value")
+        expanded = os.path.expandvars(values)
+        setattr(namespace, self.dest, str(Path(expanded).expanduser()))
+
+
+def _build_request(args: argparse.Namespace) -> Request:
+    match args.subcommand:
+        case "qpe":
+            return QPERequest(product=args.product, interval=args.interval)
+        case "qpf":
+            return QPFRequest(interval=args.interval, cycle_hour=args.cycle)
+        case "hrrr":
+            return HRRRRequest(
+                cycle_hour=args.cycle,
+                forecast_hours=tuple(args.fct_hour),
+                bbox=BoundingBox(
+                    left_lon=args.llon,
+                    right_lon=args.rlon,
+                    top_lat=args.tlat,
+                    bottom_lat=args.blat,
+                ),
+            )
+        case _:
+            raise ValueError(f"unknown subcommand: {args.subcommand}")
+
+
+def _summary(
+    result: DownloadResult,
+    *,
+    dry_run: bool,
+    subcommand: str,
+    output_dir: Path,
+) -> dict[str, Any]:
+    return {
+        "subcommand": subcommand,
+        "dry_run": dry_run,
+        "planned": len(result.plan.items),
+        "downloaded": len(result.succeeded),
+        "skipped": len(result.skipped),
+        "failed": len(result.failed),
+        "output_dir": str(output_dir),
+        "files": [item.filename for item in result.plan.items],
+    }
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="qpkit: NOAA precipitation GRIB2 downloader")
+    parser.add_argument("--version", action="version", version=f"qpkit {_qpkit_version()}")
+
+    common = argparse.ArgumentParser(add_help=False)
+    common.add_argument(
+        "-o", "--output-dir", required=True, action=PathExpandAction, dest="output_dir"
+    )
+    common.add_argument("-w", "--working-dir", action=PathExpandAction, dest="working_dir")
+    common.add_argument("--force", action="store_true")
+    common.add_argument("--dry-run", action="store_true", dest="dry_run")
+    common.add_argument("--json", action="store_true", dest="json_output")
+    common.add_argument("-l", "--log-file", action=PathExpandAction, dest="log_file")
+    common.add_argument(
+        "-n", "--log-level", type=int, choices=range(6), default=2, dest="log_level"
+    )
+
+    sub = parser.add_subparsers(dest="subcommand")
+
+    qpe = sub.add_parser("qpe", parents=[common])
+    qpe.add_argument(
+        "-p",
+        "--product",
+        choices=["MultiSensor_Pass1", "MultiSensor_Pass2", "RadarOnly"],
+        default="MultiSensor_Pass1",
+    )
+    qpe.add_argument("-i", "--interval", type=int, choices=[1, 3, 6, 12, 24, 48, 72], default=1)
+
+    utc_hour = dt.datetime.now(dt.timezone.utc).hour
+
+    qpf = sub.add_parser("qpf", parents=[common])
+    qpf.add_argument("-i", "--interval", type=int, choices=[6, 24, 48, 120], default=6)
+    qpf.add_argument("-c", "--cycle", type=int, default=utc_hour)
+
+    hrrr = sub.add_parser("hrrr", parents=[common])
+    hrrr.add_argument("-c", "--cycle", type=int, default=utc_hour)
+    hrrr.add_argument(
+        "-f",
+        "--fct-hour",
+        type=parse_forecast_hours,
+        default=parse_forecast_hours("0-18"),
+        dest="fct_hour",
+    )
+    hrrr.add_argument("--left-lon", type=float, default=0.0, dest="llon")
+    hrrr.add_argument("--right-lon", type=float, default=360.0, dest="rlon")
+    hrrr.add_argument("--top-lat", type=float, default=90.0, dest="tlat")
+    hrrr.add_argument("--bottom-lat", type=float, default=-90.0, dest="blat")
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.subcommand not in {"qpe", "qpf", "hrrr"}:
+        parser.print_help()
+        return 2
+
+    logger = build_logger(log_file=args.log_file, log_level=args.log_level)
+
+    if args.working_dir:
+        os.chdir(args.working_dir)
+
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    if not output_dir.is_dir():
+        logger.error("Output directory does not exist: %s", output_dir)
+        return 1
+
+    try:
+        request = _build_request(args)
+    except ValidationError as exc:
+        if args.json_output:
+            print(json.dumps({"error": str(exc)}))
+        else:
+            logger.error("Invalid request: %s", exc)
+        return 1
+
+    try:
+        with QPKit(logger=logger) as kit:
+            plan = kit.build_plan(request, output_dir, force=args.force)
+            if args.dry_run:
+                result = DownloadResult(plan=plan, skipped=plan.skipped)
+            else:
+                result = kit.downloader.download(plan)
+
+            summary = _summary(
+                result,
+                dry_run=args.dry_run,
+                subcommand=args.subcommand,
+                output_dir=output_dir,
+            )
+
+            if args.json_output:
+                print(json.dumps(summary, sort_keys=True))
+            else:
+                logger.info(
+                    "Summary: planned=%s downloaded=%s skipped=%s failed=%s",
+                    summary["planned"],
+                    summary["downloaded"],
+                    summary["skipped"],
+                    summary["failed"],
+                )
+                if args.dry_run:
+                    logger.info("Dry run: no files were downloaded.")
+
+            return 1 if summary["failed"] > 0 else 0
+    except QPKitError as exc:
+        if args.json_output:
+            print(json.dumps({"error": str(exc)}))
+        else:
+            logger.error(str(exc))
+        return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

qpkit/downloader.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import logging
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+import httpx
+
+from .http import HttpClientProtocol
+from .models import DownloadItem, DownloadPlan, DownloadResult
+
+
+class Downloader:
+    def __init__(
+        self,
+        http: HttpClientProtocol,
+        logger: logging.Logger,
+        *,
+        max_workers: int = 4,
+        max_retries: int = 3,
+        backoff_seconds: float = 1.0,
+    ) -> None:
+        self.http = http
+        self.logger = logger
+        self.max_workers = max_workers
+        self.max_retries = max_retries
+        self.backoff_seconds = backoff_seconds
+
+    def _download_one(self, item: DownloadItem, output_path: Path) -> Path | None:
+        temp_path = output_path.with_name(f"{output_path.name}.part")
+        temp_path.unlink(missing_ok=True)
+
+        for attempt in range(self.max_retries + 1):
+            try:
+                self.http.download(item.url, temp_path)
+                temp_path.replace(output_path)
+                self.logger.info("Download saved to: %s", output_path)
+                return output_path
+            except (httpx.HTTPError, OSError) as exc:
+                temp_path.unlink(missing_ok=True)
+                if attempt < self.max_retries:
+                    delay = self.backoff_seconds * (2**attempt)
+                    self.logger.warning(
+                        "Download retry %s/%s for %s in %.1fs: %s",
+                        attempt + 1,
+                        self.max_retries,
+                        output_path.name,
+                        delay,
+                        exc,
+                    )
+                    time.sleep(delay)
+                    continue
+
+                self.logger.warning(
+                    "Failed download %s after retries: %s", output_path.name, exc
+                )
+                return None
+
+        return None
+
+    def download(self, plan: DownloadPlan) -> DownloadResult:
+        if not plan.items:
+            return DownloadResult(plan=plan, skipped=plan.skipped)
+
+        output_dir = plan.output_dir
+        workers = min(self.max_workers, len(plan.items))
+        succeeded: list[Path] = []
+        failed: list[str] = []
+
+        with ThreadPoolExecutor(max_workers=workers) as pool:
+            future_to_item = {
+                pool.submit(self._download_one, item, output_dir / item.filename): item
+                for item in plan.items
+            }
+            for future in as_completed(future_to_item):
+                item = future_to_item[future]
+                result = future.result()
+                if result is None:
+                    failed.append(item.url)
+                else:
+                    succeeded.append(result)
+
+        return DownloadResult(
+            succeeded=tuple(succeeded),
+            failed=tuple(failed),
+            skipped=plan.skipped,
+            plan=plan,
+        )