internship/swash/processing/swash.py

import argparse
import configparser
import logging
import os
import pathlib
import shutil
import subprocess
import sys
import tempfile

parser = argparse.ArgumentParser(description="Run swash model")
parser.add_argument("-v", "--verbose", action="count", default=0)
parser.add_argument("-c", "--config", default="config.ini")
args = parser.parse_args()

logging.basicConfig(level=max((10, 20 - 10 * args.verbose)))
log = logging.getLogger("swash")

log.info("Starting swash model")
config = configparser.ConfigParser()
config.read(args.config)

inp = pathlib.Path(config.get("swash", "input"))
out = pathlib.Path(config.get("swash", "out"))
if out.exists():
    log.error(f"Swash output '{out}' already exists")
    sys.exit(1)
out.parent.mkdir(parents=True, exist_ok=True)

with tempfile.TemporaryDirectory(prefix="swash_", dir=".") as tmp_raw:
    tmpdir = pathlib.Path(tmp_raw)

    log.info(f"Copying files to '{tmpdir}'")
    shutil.copy2(inp, tmpdir)
    if config.getboolean("swash", "nb", fallback=False):
        path = pathlib.Path(config.get("data", "out_nb"))
    else:
        path = pathlib.Path(config.get("data", "out"))
    shutil.copytree(path, tmpdir, dirs_exist_ok=True)

    if config.has_option("swash", "mpi"):
        mpi = ("-mpi", config.get("swash", "mpi"))
        log.info(f"Using mpi with {mpi}")
    else:
        mpi = ()

    with open(tmpdir.joinpath("sws.log"), "w") as logfile:
        log.info(f"Runing swash in '{tmpdir}'")
        path = pathlib.Path(config.get("swash", "path"))

        cmd = (path.joinpath("swashrun"), "-input", inp.name, *mpi)
        log.info(f"Running {cmd}")

        swash_run = subprocess.Popen(
            cmd,
            cwd=tmpdir,
            stdout=logfile,
            stderr=logfile,
            env={"PATH": f"{path}:{os.environ['PATH']}"},
        )

        code = swash_run.wait()
        if code != 0:
            log.error(f"Swash returned error code {code}")

    log.info(f"Moving swash output to '{out}'")
    shutil.move(tmpdir, out)

for f in out.rglob("*"):
    if f.is_file():
        f.chmod(0o444)

log.info(f"Swash model finished successfully")