feat(dfn): toml load support, switch to tomli, add tests (MODFLOW-USG…

…S#173)

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,13 +23,43 @@ 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,
     ):
         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)

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}")

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)

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"},