wip

Author: wpbonelli

autotest/test_dfn.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from modflow_devtools.dfn import _load_common, load, load_all, load_tree
+from modflow_devtools.dfn import _load_common, load, load_all
 from modflow_devtools.dfn.fetch import fetch_dfns
 from modflow_devtools.dfn2toml import convert
 from modflow_devtools.markers import requires_pkg
@@ -29,14 +29,11 @@ def pytest_generate_tests(metafunc):
 
     if "toml_name" in metafunc.fixturenames:
         convert(DFN_DIR, TOML_DIR)
-        dfn_paths = list(DFN_DIR.glob("*.dfn"))
         expected_toml_paths = [
-            TOML_DIR / f"{dfn.stem.replace('-nam', '')}.toml"
-            for dfn in dfn_paths
-            if "common" not in dfn.stem
+            dfn for dfn in DFN_DIR.glob("*.dfn") if "common" not in dfn.stem
         ]
-        toml_names = [toml.stem for toml in TOML_DIR.glob("*.toml")]
         assert all(toml_path.exists() for toml_path in expected_toml_paths)
+        toml_names = [toml.stem for toml in TOML_DIR.glob("*.toml")]
         metafunc.parametrize("toml_name", toml_names, ids=toml_names)
 
 
@@ -47,7 +44,7 @@ def test_load_v1(dfn_name):
         (DFN_DIR / f"{dfn_name}.dfn").open() as dfn_file,
     ):
         common, _ = _load_common(common_file)
-        dfn = load(dfn_file, name=dfn_name, common=common)
+        dfn = load(dfn_file, name=dfn_name, format="dfn", common=common)
         assert any(dfn.fields)
 
 
@@ -61,65 +58,57 @@ def test_load_v2(toml_name):
 @requires_pkg("boltons")
 @pytest.mark.parametrize("schema_version", list(SPEC_DIRS.keys()))
 def test_load_all(schema_version):
-    path = SPEC_DIRS[schema_version]
-    dfns = load_all(path)
+    dfns = load_all(path=SPEC_DIRS[schema_version])
     assert all(any(dfn.fields) for dfn in dfns.values())
 
 
-@requires_pkg("boltons")
-def test_load_tree():
-    import tempfile
-
+@requires_pkg("boltons", "tomli")
+def test_convert(function_tmpdir):
     import tomli
 
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        tmp_path = Path(tmp_dir)
-        convert(DFN_DIR, tmp_path)
-
-        # Test file conversion and naming
-        assert (tmp_path / "sim.toml").exists()
-        assert (tmp_path / "gwf.toml").exists()
-        assert not (tmp_path / "sim-nam.toml").exists()
-
-        # Test parent relationships in files
-        with (tmp_path / "sim.toml").open("rb") as f:
-            sim_data = tomli.load(f)
-        assert sim_data["name"] == "sim"
-        assert "parent" not in sim_data
-
-        with (tmp_path / "gwf.toml").open("rb") as f:
-            gwf_data = tomli.load(f)
-        assert gwf_data["name"] == "gwf"
-        assert gwf_data["parent"] == "sim"
-
-        dfns = load_all(tmp_path)
-        root = load_tree(tmp_path)
-        roots = []
-        for dfn in dfns.values():
-            if dfn.parent:
-                assert dfn.parent in dfns
-            else:
-                roots.append(dfn.name)
-        assert len(roots) == 1
-        assert root.name == "sim"
-        assert root == roots[0]
-
-        model_types = ["gwf", "gwt", "gwe"]
-        models = root.children or {}
-        for model_type in model_types:
-            if model_type in models:
-                assert models[model_type].name == model_type
-                assert models[model_type].parent == "sim"
-
-        if "gwf" in models:
-            pkgs = models["gwf"].children or {}
-            gwf_packages = [
-                k for k in pkgs if k.startswith("gwf-") and isinstance(pkgs[k], dict)
-            ]
-            assert len(gwf_packages) > 0
-
-            if dis := pkgs.get("gwf-dis", None):
-                assert dis.name == "gwf-dis"
-                assert dis.parent == "gwf"
-                assert "options" in (dis.blocks or {})
-                assert "dimensions" in (dis.blocks or {})
+    convert(DFN_DIR, function_tmpdir)
+
+    assert (function_tmpdir / "sim-nam.toml").exists()
+    assert (function_tmpdir / "gwf-nam.toml").exists()
+
+    with (function_tmpdir / "sim-nam.toml").open("rb") as f:
+        sim_data = tomli.load(f)
+    assert sim_data["name"] == "sim-nam"
+    assert sim_data["schema_version"] == "2"
+    assert "parent" not in sim_data
+
+    with (function_tmpdir / "gwf-nam.toml").open("rb") as f:
+        gwf_data = tomli.load(f)
+    assert gwf_data["name"] == "gwf-nam"
+    assert gwf_data["parent"] == "sim-nam"
+    assert gwf_data["schema_version"] == "2"
+
+    dfns = load_all(function_tmpdir)
+    roots = []
+    for dfn in dfns.values():
+        if dfn.parent:
+            assert dfn.parent in dfns
+        else:
+            roots.append(dfn.name)
+    assert len(roots) == 1
+    root = dfns[roots[0]]
+    assert root.name == "sim-nam"
+
+    models = root.children or {}
+    for mdl in models:
+        assert models[mdl].name == mdl
+        assert models[mdl].parent == "sim-nam"
+
+    if gwf := models.get("gwf-nam", None):
+        pkgs = gwf.children or {}
+        pkgs = {
+            k: v
+            for k, v in pkgs.items()
+            if k.startswith("gwf-") and isinstance(v, dict)
+        }
+        assert len(pkgs) > 0
+        if dis := pkgs.get("gwf-dis", None):
+            assert dis.name == "gwf-dis"
+            assert dis.parent == "gwf"
+            assert "options" in (dis.blocks or {})
+            assert "dimensions" in (dis.blocks or {})

modflow_devtools/dfn/__init__.py

@@ -29,7 +29,7 @@
 from modflow_devtools.dfn.schema.ref import Ref
 from modflow_devtools.dfn.schema.v1 import FieldV1
 from modflow_devtools.dfn.schema.v2 import FieldV2
-from modflow_devtools.misc import try_literal_eval
+from modflow_devtools.misc import drop_none_or_empty, try_literal_eval
 
 __all__ = [
     "SCALAR_TYPES",
@@ -41,6 +41,7 @@
     "FieldV1",
     "Fields",
     "Ref",
+    "flatten",
     "infer_tree",
     "load",
     "load_all",
@@ -62,8 +63,8 @@ class Dfn:
     MODFLOW 6 input component definition.
     """
 
-    name: str
     schema_version: Version
+    name: str
     parent: str | None = None
     advanced: bool = False
     multi: bool = False
@@ -313,13 +314,15 @@ def remove_attrs(path, key, value):
     def map(self, dfn: Dfn) -> Dfn:
         if dfn.schema_version == (v2 := Version("2")):
             return dfn
+
         return Dfn(
             name=dfn.name,
             advanced=dfn.advanced,
             multi=dfn.multi,
             ref=dfn.ref,
             blocks=MapV1To2.map_blocks(dfn),
             schema_version=v2,
+            parent=dfn.parent,
         )
 
 
@@ -396,63 +399,73 @@ def infer_tree(dfns: Dfns) -> Dfn:
     Assumes all DFNs are of the same schema version.
     """
 
-    def drop_none_or_empty(path, key, value):
-        if value is None or value == "" or value == [] or value == {}:
-            return False
-        return True
-
-    def add_parent(dfn):
-        dfn = dfn.copy()
-        dfn_name = dfn["name"]
-        if dfn_name == "sim-nam":
-            # simulation name file
-            dfn.name = "sim"
+    def set_parent(dfn):
+        dfn = asdict(dfn)
+        if (dfn_name := dfn["name"]) == "sim-nam":
+            pass
         elif dfn_name.endswith("-nam"):
-            # model name files
-            model_type = dfn_name[:-4]  # Remove "-nam"
-            dfn.name = model_type
-            dfn.parent = "sim"
+            dfn["parent"] = "sim-nam"
         elif (
             dfn_name.startswith("exg-")
             or dfn_name.startswith("sln-")
             or dfn_name.startswith("utl-")
         ):
-            # exchanges, solutions, standalone utilities
-            dfn.parent = "sim"
+            dfn["parent"] = "sim-nam"
         elif "-" in dfn_name:
-            # packages
-            model_type = dfn_name.split("-")[0]
-            dfn.parent = model_type
-
-        return remap(dfn, visit=drop_none_or_empty)
-
-    dfns = {name: add_parent(dfn) for name, dfn in dfns.items()}
-    first = next(iter(dfns.values()), None)
-    schema_version = first.schema_version if first else Version("1")
-
-    if schema_version == Version("1"):
-        raise NotImplementedError("Structure inference from v1 schema not implemented")
-    elif schema_version == Version("2"):
-        if (
-            len(
-                roots := [
-                    (name, dfn) for name, dfn in dfns.items() if dfn.parent is None
-                ]
+            mdl = dfn_name.split("-")[0]
+            dfn["parent"] = f"{mdl}-nam"
+
+        return Dfn(**remap(dfn, visit=drop_none_or_empty))
+
+    dfns = {name: set_parent(dfn) for name, dfn in dfns.items()}
+    first_dfn = next(iter(dfns.values()), None)
+    match schema_version := str(
+        first_dfn.schema_version if first_dfn else Version("1")
+    ):
+        case "1":
+            raise NotImplementedError("Tree inference from v1 schema not implemented")
+        case "2":
+            if (
+                nroots := len(
+                    roots := [
+                        (name, dfn) for name, dfn in dfns.items() if dfn.parent is None
+                    ]
+                )
+            ) != 1:
+                raise ValueError(f"Expected one root component, found {nroots}")
+
+            def get_children(node_name: str) -> Dfns:
+                return {
+                    name: dfn for name, dfn in dfns.items() if dfn.parent == node_name
+                }
+
+            return Dfn(
+                name=(root_name := roots[0][0]),
+                blocks=dfns[root_name].blocks,
+                children=get_children(root_name),
+                schema_version=Version(schema_version),
+            )
+        case _:
+            raise ValueError(
+                f"Unsupported schema version: {schema_version}. Expected 1 or 2."
             )
-            != 1
-        ):
-            raise ValueError(f"Expected one root component, found {len(roots)}")
 
-        def get_children(node_name: str) -> Dfns:
-            return {name: dfn for name, dfn in dfns.items() if dfn.parent == node_name}
 
-        return Dfn(
-            name=(root_name := roots[0][0]),
-            blocks=dfns[root_name].blocks,
-            children=get_children(root_name),
-            schema_version=schema_version,
-        )
-    raise ValueError(f"Unsupported schema version: {schema_version}. Expected 1 or 2.")
+def flatten(dfn: Dfn) -> Dfns:
+    """
+    Flatten a MODFLOW 6 input component hierarchy to a flat spec:
+    unlinked DFNs, i.e. without `children` populated, only `parent`.
+
+    Returns a dictionary of all components in the tree.
+    """
+
+    def _flatten(dfn: Dfn) -> Dfns:
+        dfns = {dfn.name: replace(dfn, children=None)}
+        for child in (dfn.children or {}).values():
+            dfns.update(_flatten(child))
+        return dfns
+
+    return _flatten(dfn)
 
 
 def load_tree(dfndir: str | PathLike) -> Dfn:

modflow_devtools/dfn2toml.py

@@ -8,7 +8,7 @@
 import tomli_w as tomli
 from boltons.iterutils import remap
 
-from modflow_devtools.dfn import load_all, map
+from modflow_devtools.dfn import flatten, load_tree, map
 from modflow_devtools.misc import drop_none_or_empty
 
 # mypy: ignore-errors
@@ -18,7 +18,7 @@ def convert(indir: PathLike, outdir: PathLike, schema_version: str = "2") -> Non
     indir = Path(indir).expanduser().absolute()
     outdir = Path(outdir).expanduser().absolute()
     outdir.mkdir(exist_ok=True, parents=True)
-    dfns_in = load_all(indir)
+    dfns_in = flatten(load_tree(indir))
     dfns = {
         name: map(dfn, schema_version=schema_version) for name, dfn in dfns_in.items()
     }

modflow_devtools/misc.py

@@ -4,6 +4,7 @@
 import traceback
 from _warnings import warn
 from ast import literal_eval
+from collections.abc import Iterable
 from contextlib import contextmanager
 from enum import Enum
 from functools import wraps
@@ -585,9 +586,9 @@ def try_literal_eval(value: str) -> Any:
 
 def drop_none_or_empty(path, key, value):
     """
-    Drop dictionary items with None or empty string values.
+    Drop dictionary items with None or empty values.
     For use with `boltons.iterutils.remap`.
     """
-    if value is None or value == "":
+    if value is None or (isinstance(value, Iterable) and not any(value)):
         return False
     return True