Managing YAML Configs for Geospatial CLI Workflows

pip install "pyyaml>=6.0" "pydantic>=2.0" "click>=8.1"

# pipeline.yaml
workspace: /data/rasters
input_glob: "**/*.tif"
output_dir: /data/output

gdal:
  GDAL_CACHEMAX: 512
  GDAL_NUM_THREADS: 8
  OGR_ENABLE_PARTIAL_REPROJECTION: true

spatial:
  src_crs: "EPSG:4326"
  dst_crs: "EPSG:3857"
  resampling: bilinear
  tile_size: 512

#!/usr/bin/env python3
"""geo_pipeline.py — validated YAML config loader for geospatial batch CLIs."""
import glob
import os
from pathlib import Path
from typing import List, Optional

import click
import yaml
from pydantic import BaseModel, Field, ValidationError, field_validator


# ── Schema models ─────────────────────────────────────────────────────────────

class GdalEnv(BaseModel):
    """GDAL/OGR runtime settings injected via os.environ before any I/O."""
    GDAL_CACHEMAX: int = Field(default=256, ge=64, le=4096,
                               description="Raster block cache in MB")
    GDAL_NUM_THREADS: int = Field(default=4, ge=1, le=64)
    OGR_ENABLE_PARTIAL_REPROJECTION: bool = Field(default=True)


class SpatialParams(BaseModel):
    src_crs: str = Field(default="EPSG:4326")
    dst_crs: str = Field(default="EPSG:3857")
    resampling: str = Field(default="bilinear")
    tile_size: int = Field(default=512, ge=128, le=4096)

    @field_validator("src_crs", "dst_crs")
    @classmethod
    def validate_crs(cls, v: str) -> str:
        if not (v.upper().startswith("EPSG:") or v.upper().startswith("PROJ:")):
            raise ValueError(f"CRS must begin with EPSG: or PROJ:, got {v!r}")
        return v

    @field_validator("resampling")
    @classmethod
    def validate_resampling(cls, v: str) -> str:
        valid = {"nearest", "bilinear", "cubic", "cubicspline", "lanczos",
                 "average", "mode"}
        if v.lower() not in valid:
            raise ValueError(f"resampling must be one of {sorted(valid)}")
        return v.lower()


class BatchConfig(BaseModel):
    workspace: Path
    input_glob: str
    output_dir: Path
    gdal: GdalEnv = Field(default_factory=GdalEnv)
    spatial: SpatialParams = Field(default_factory=SpatialParams)

    @field_validator("workspace", "output_dir", mode="before")
    @classmethod
    def resolve_paths(cls, v: object) -> Path:
        return Path(str(v)).expanduser().resolve()

    def apply_gdal_env(self) -> None:
        """Inject GDAL settings into os.environ before any rasterio or GDAL call."""
        for key, value in self.gdal.model_dump().items():
            os.environ[key] = str(value)

    def resolve_inputs(self) -> List[Path]:
        """Expand input_glob relative to workspace into a sorted, deduplicated list."""
        pattern = str(self.workspace / self.input_glob)
        return sorted({Path(p) for p in glob.glob(pattern, recursive=True)})


# ── CLI entry-point ────────────────────────────────────────────────────────────

@click.command()
@click.option(
    "--config", "cfg_path",
    type=click.Path(exists=True, path_type=Path),
    required=True,
    help="Path to the YAML pipeline config file.",
)
@click.option(
    "--threads", type=int, default=None,
    help="Override GDAL_NUM_THREADS from the config file.",
)
@click.option(
    "--dst-crs", "dst_crs", default=None,
    help="Override spatial.dst_crs (e.g. EPSG:32633).",
)
def run_pipeline(cfg_path: Path, threads: Optional[int], dst_crs: Optional[str]) -> None:
    """Load, validate, and execute a geospatial batch pipeline from a YAML config."""
    with cfg_path.open("r") as fh:
        raw: dict = yaml.safe_load(fh)

    # ① Merge CLI overrides into raw dict BEFORE constructing the model.
    #    This is the key pattern: the model never needs to know about Click.
    if threads is not None:
        raw.setdefault("gdal", {})["GDAL_NUM_THREADS"] = threads
    if dst_crs is not None:
        raw.setdefault("spatial", {})["dst_crs"] = dst_crs

    # ② Validate everything in one shot — fail fast with a clear message.
    try:
        config = BatchConfig(**raw)
    except ValidationError as exc:
        click.echo(f"Config validation failed:\n{exc}", err=True)
        raise click.exceptions.Exit(code=2)

    # ③ Inject GDAL env before any I/O (rasterio, pyogrio, osgeo.gdal).
    config.apply_gdal_env()

    # ④ Audit log: print the resolved config so CI logs are self-documenting.
    click.echo(config.model_dump_json(indent=2))

    # ⑤ Resolve input files and abort early if the glob matches nothing.
    inputs = config.resolve_inputs()
    if not inputs:
        click.echo(
            f"No files matched '{config.input_glob}' under {config.workspace}",
            err=True,
        )
        raise click.exceptions.Exit(code=2)

    click.echo(
        f"Processing {len(inputs)} files: "
        f"{config.spatial.src_crs} → {config.spatial.dst_crs} "
        f"@ {config.spatial.resampling} resampling"
    )
    # ⑥ Pipeline execution continues here with fully validated config.


if __name__ == "__main__":
    run_pipeline()

# WRONG — GDAL_CACHEMAX is ignored for this handle
with rasterio.open(src) as ds:
    config.apply_gdal_env()   # too late; cache is already allocated
    data = ds.read()

# CORRECT — inject before any open() call
config.apply_gdal_env()
with rasterio.open(src) as ds:
    data = ds.read()

# 1. Confirm GDAL env variables were injected (visible in process env)
python - <<'EOF'
import os, yaml, pydantic
# Quick smoke-test: construct config and check env after apply
from geo_pipeline import BatchConfig
import yaml
raw = yaml.safe_load(open("pipeline.yaml"))
cfg = BatchConfig(**raw)
cfg.apply_gdal_env()
assert os.environ["GDAL_CACHEMAX"] == str(cfg.gdal.GDAL_CACHEMAX)
print("GDAL env OK:", os.environ["GDAL_CACHEMAX"], "MB")
EOF

# 2. Check that an output GeoTIFF is in the expected CRS (requires GDAL CLI tools)
gdalinfo output/tile_0001.tif | grep -E "EPSG|CoordSys"