Layering TOML and Env Config for Raster Pipelines

#!/usr/bin/env python3
"""
Layered configuration for a raster pipeline CLI.
Order (lowest to highest): defaults -> TOML -> env -> CLI flags.
Requires Python 3.11+ for tomllib.
"""
from __future__ import annotations

import os
import tomllib
from dataclasses import dataclass, fields, replace
from pathlib import Path
from typing import Any


ENV_PREFIX = "MYTOOL_"
TOML_TABLE = ("tool", "mytool")


@dataclass(frozen=True)
class RasterConfig:
    target_epsg: int = 4326          # output CRS as a bare EPSG code
    max_workers: int = 4             # parallel warp workers
    chunk_size: int = 512            # tile edge in pixels for windowed reads
    output_dir: Path = Path("./out")  # where reprojected rasters land
    overwrite: bool = False          # replace existing outputs


def _coerce(field_type: Any, raw: str) -> Any:
    """Coerce a raw string (from env) to the dataclass field's type.

    Environment values are ALWAYS strings; TOML values are already typed.
    bool needs special handling because bool("0") is True in Python.
    """
    if field_type is bool:
        return raw.strip().lower() in {"1", "true", "yes", "on"}
    if field_type is int:
        return int(raw)
    if field_type is Path:
        return Path(raw)
    return raw  # str fields pass through unchanged


def _load_toml(path: Path) -> dict[str, Any]:
    """Return the [tool.mytool] table, or {} if the file/table is absent."""
    if not path.is_file():
        return {}
    with path.open("rb") as fh:          # tomllib requires binary mode
        doc = tomllib.load(fh)
    table: Any = doc
    for key in TOML_TABLE:
        table = table.get(key, {}) if isinstance(table, dict) else {}
    return table if isinstance(table, dict) else {}


def _load_env() -> dict[str, str]:
    """Collect MYTOOL_-prefixed vars, lowercased to field names."""
    out: dict[str, str] = {}
    for key, value in os.environ.items():
        if key.startswith(ENV_PREFIX):
            out[key[len(ENV_PREFIX):].lower()] = value
    return out


def load_config(
    toml_path: Path = Path("pyproject.toml"),
    cli_overrides: dict[str, Any] | None = None,
) -> tuple[RasterConfig, dict[str, str]]:
    """Resolve config across all four layers.

    Returns the final RasterConfig plus a provenance map recording which
    layer supplied each field's winning value.
    """
    field_types = {f.name: f.type for f in fields(RasterConfig)}
    valid = set(field_types)

    config = RasterConfig()                      # layer 1: defaults
    provenance = {name: "default" for name in valid}

    # layer 2: TOML — values are already correctly typed by tomllib
    for name, value in _load_toml(toml_path).items():
        if name in valid:
            coerced = Path(value) if field_types[name] is Path else value
            config = replace(config, **{name: coerced})
            provenance[name] = "toml"

    # layer 3: env — every value is a string, so coerce per field type
    for name, raw in _load_env().items():
        if name in valid:
            config = replace(config, **{name: _coerce(field_types[name], raw)})
            provenance[name] = "env"

    # layer 4: CLI flags — merge ONLY keys the user actually set (not None)
    for name, value in (cli_overrides or {}).items():
        if name in valid and value is not None:
            config = replace(config, **{name: value})
            provenance[name] = "flag"

    return config, provenance


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Raster pipeline runner")
    parser.add_argument("--config", type=Path, default=Path("pyproject.toml"))
    # CLI flags DEFAULT TO None so an unset flag never clobbers lower layers.
    parser.add_argument("--target-epsg", type=int, default=None)
    parser.add_argument("--max-workers", type=int, default=None)
    parser.add_argument("--chunk-size", type=int, default=None)
    parser.add_argument("--output-dir", type=Path, default=None)
    parser.add_argument("--overwrite", action="store_true", default=None)
    args = parser.parse_args()

    overrides = {
        "target_epsg": args.target_epsg,
        "max_workers": args.max_workers,
        "chunk_size": args.chunk_size,
        "output_dir": args.output_dir,
        "overwrite": args.overwrite,
    }
    cfg, prov = load_config(args.config, overrides)

    print("Resolved raster pipeline configuration:")
    for f in fields(cfg):
        print(f"  {f.name:<12} = {getattr(cfg, f.name)!r:<20} (from {prov[f.name]})")

[tool.mytool]
target_epsg = 3857
max_workers = 8
output_dir = "./reprojected"

# TOML sets max_workers=8; env overrides target_epsg; a flag overrides chunk_size.
export MYTOOL_TARGET_EPSG=32633
python config_loader.py --chunk-size 1024

# Expected output:
#   target_epsg  = 32633                (from env)
#   max_workers  = 8                    (from toml)
#   chunk_size   = 1024                 (from flag)
#   output_dir   = PosixPath('reprojected')  (from toml)
#   overwrite    = False                (from default)

import os
from pathlib import Path
from config_loader import load_config

os.environ["MYTOOL_TARGET_EPSG"] = "32633"
cfg, prov = load_config(
    toml_path=Path("pyproject.toml"),
    cli_overrides={"chunk_size": 1024, "max_workers": None},
)

assert cfg.target_epsg == 32633 and prov["target_epsg"] == "env"
assert cfg.chunk_size == 1024 and prov["chunk_size"] == "flag"
assert prov["max_workers"] == "toml"      # unset flag did NOT overwrite
assert isinstance(cfg.max_workers, int)   # coercion held
print("precedence chain verified")