From 01b95b0745a4fc4767cf8074738e2968f1e206b1 Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Tue, 14 Jan 2025 16:00:40 -0500 Subject: [PATCH] feat(dfn): toml load support, switch to tomli, add tests (#173) --- autotest/test_dfn.py | 34 +++++++++++++++++++++++++++++++++- modflow_devtools/dfn.py | 36 ++++++++++++++++++++++++++++++++++-- modflow_devtools/dfn2toml.py | 25 +++++++++++++++---------- pyproject.toml | 6 ++++-- 4 files changed, 86 insertions(+), 15 deletions(-) diff --git a/autotest/test_dfn.py b/autotest/test_dfn.py index 47b122c..a80d0fc 100644 --- a/autotest/test_dfn.py +++ b/autotest/test_dfn.py @@ -1,10 +1,12 @@ from pathlib import Path from modflow_devtools.dfn import Dfn, get_dfns +from modflow_devtools.dfn2toml import convert from modflow_devtools.markers import requires_pkg PROJ_ROOT = Path(__file__).parents[1] DFN_PATH = PROJ_ROOT / "autotest" / "temp" / "dfn" +TOML_PATH = DFN_PATH / "toml" MF6_OWNER = "MODFLOW-USGS" MF6_REPO = "modflow6" MF6_REF = "develop" @@ -21,9 +23,20 @@ def pytest_generate_tests(metafunc): ] metafunc.parametrize("dfn_name", dfn_names, ids=dfn_names) + if "toml_name" in metafunc.fixturenames: + convert(DFN_PATH, TOML_PATH) + dfns = list(DFN_PATH.glob("*.dfn")) + assert all( + (TOML_PATH / f"{dfn.stem}.toml").is_file() + for dfn in dfns + if "common" not in dfn.stem + ) + toml_names = [toml.stem for toml in TOML_PATH.glob("*.toml")] + metafunc.parametrize("toml_name", toml_names, ids=toml_names) + @requires_pkg("boltons") -def test_dfn_load(dfn_name): +def test_load_v1(dfn_name): with ( (DFN_PATH / "common.dfn").open() as common_file, (DFN_PATH / f"{dfn_name}.dfn").open() as dfn_file, @@ -31,3 +44,22 @@ def test_dfn_load(dfn_name): common, _ = Dfn._load_v1_flat(common_file) dfn = Dfn.load(dfn_file, name=dfn_name, common=common) assert any(dfn) + + +@requires_pkg("boltons") +def test_load_all_v1(): + dfns = Dfn.load_all(DFN_PATH) + assert any(dfns) + + +@requires_pkg("boltons") +def test_load_v2(toml_name): + with (TOML_PATH / f"{toml_name}.toml").open(mode="rb") as toml_file: + toml = Dfn.load(toml_file, name=toml_name, version=2) + assert any(toml) + + +@requires_pkg("boltons") +def test_load_all_v2(): + toml = Dfn.load_all(TOML_PATH, version=2) + assert any(toml) diff --git a/modflow_devtools/dfn.py b/modflow_devtools/dfn.py index cbbe94c..79b1404 100644 --- a/modflow_devtools/dfn.py +++ b/modflow_devtools/dfn.py @@ -14,8 +14,11 @@ ) from warnings import warn +import tomli from boltons.dictutils import OMD +from modflow_devtools.download import download_and_unzip + # DFN representation with a # parser for the DFN format @@ -466,6 +469,17 @@ def _multi() -> bool: **blocks, ) + @classmethod + def _load_v2(cls, f, name) -> "Dfn": + # load data + data = tomli.load(f) + + # if name provided, make sure it matches + if name and name != data.get("name", None): + raise ValueError(f"Name mismatch, expected {name}") + + return cls(**data) + @classmethod def load( cls, @@ -480,6 +494,8 @@ def load( if version == 1: return cls._load_v1(f, name, **kwargs) + elif version == 2: + return cls._load_v2(f, name) else: raise ValueError(f"Unsupported version, expected one of {version.__args__}") @@ -516,12 +532,30 @@ def _load_all_v1(dfndir: PathLike) -> Dfns: return dfns + @staticmethod + def _load_all_v2(dfndir: PathLike) -> Dfns: + # find definition files + paths: list[Path] = [ + p for p in dfndir.glob("*.toml") if p.stem not in ["common", "flopy"] + ] + + # load all the input definitions + dfns: Dfns = {} + for path in paths: + with path.open(mode="rb") as f: + dfn = Dfn.load(f, name=path.stem, version=2) + dfns[path.stem] = dfn + + return dfns + @staticmethod def load_all(dfndir: PathLike, version: DfnFmtVersion = 1) -> Dfns: """Load all input definitions from the given directory.""" if version == 1: return Dfn._load_all_v1(dfndir) + elif version == 2: + return Dfn._load_all_v2(dfndir) else: raise ValueError(f"Unsupported version, expected one of {version.__args__}") @@ -532,8 +566,6 @@ def load_all(dfndir: PathLike, version: DfnFmtVersion = 1) -> Dfns: def get_dfns( owner: str, repo: str, ref: str, outdir: Union[str, PathLike], verbose: bool = False ): - from modflow_devtools.download import download_and_unzip - url = f"https://github.com/{owner}/{repo}/archive/{ref}.zip" if verbose: print(f"Downloading MODFLOW 6 repository from {url}") diff --git a/modflow_devtools/dfn2toml.py b/modflow_devtools/dfn2toml.py index 2a4ef5e..9644935 100644 --- a/modflow_devtools/dfn2toml.py +++ b/modflow_devtools/dfn2toml.py @@ -1,8 +1,13 @@ import argparse from collections.abc import Mapping +from os import PathLike from pathlib import Path from typing import Any +import tomli_w as tomli + +from modflow_devtools.dfn import Dfn + class Shim: @staticmethod @@ -27,12 +32,17 @@ def apply(d: dict) -> dict: return Shim._attach_children(Shim._drop_none(d)) -if __name__ == "__main__": - """Convert DFN files to TOML.""" +def convert(indir: PathLike, outdir: PathLike): + indir = Path(indir).expanduser().absolute() + outdir = Path(outdir).expanduser().absolute() + outdir.mkdir(exist_ok=True, parents=True) + for dfn in Dfn.load_all(indir).values(): + with Path.open(outdir / f"{dfn['name']}.toml", "wb") as f: + tomli.dump(Shim.apply(dfn), f) - import tomlkit - from modflow_devtools.dfn import Dfn +if __name__ == "__main__": + """Convert DFN files to TOML.""" parser = argparse.ArgumentParser(description="Convert DFN files to TOML.") parser.add_argument( @@ -47,9 +57,4 @@ def apply(d: dict) -> dict: help="Output directory.", ) args = parser.parse_args() - indir = Path(args.indir) - outdir = Path(args.outdir) - outdir.mkdir(exist_ok=True, parents=True) - for dfn in Dfn.load_all(indir).values(): - with Path.open(outdir / f"{dfn['name']}.toml", "w") as f: - tomlkit.dump(Shim.apply(dfn), f) + convert(args.indir, args.outdir) diff --git a/pyproject.toml b/pyproject.toml index a54f341..29c2360 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -74,7 +74,8 @@ docs = [ ] dfn = [ "boltons", - "tomlkit" + "tomli", + "tomli-w" ] dev = ["modflow-devtools[lint,test,docs,dfn]"] @@ -110,7 +111,8 @@ docs = [ ] dfn = [ "boltons", - "tomlkit" + "tomli", + "tomli-w" ] dev = [ {include-group = "build"},