almost there

Author: wpbonelli

autotest/test_dfn.py

@@ -2,15 +2,15 @@
 
 import pytest
 
-from modflow_devtools.dfn import Dfn, load
+from modflow_devtools.dfn import _load_common, load, load_all, load_tree
 from modflow_devtools.dfn.fetch import fetch_dfns
 from modflow_devtools.dfn2toml import convert
 from modflow_devtools.markers import requires_pkg
 
 PROJ_ROOT = Path(__file__).parents[1]
 DFN_DIR = PROJ_ROOT / "autotest" / "temp" / "dfn"
 TOML_DIR = DFN_DIR / "toml"
-VERSIONS = {1: DFN_DIR, 2: TOML_DIR}
+SPEC_DIRS = {1: DFN_DIR, 2: TOML_DIR}
 MF6_OWNER = "MODFLOW-ORG"
 MF6_REPO = "modflow6"
 MF6_REF = "develop"
@@ -30,12 +30,13 @@ def pytest_generate_tests(metafunc):
     if "toml_name" in metafunc.fixturenames:
         convert(DFN_DIR, TOML_DIR)
         dfn_paths = list(DFN_DIR.glob("*.dfn"))
-        assert all(
-            (TOML_DIR / f"{dfn.stem.replace('-nam', '')}.toml").is_file()
+        expected_toml_paths = [
+            TOML_DIR / f"{dfn.stem.replace('-nam', '')}.toml"
             for dfn in dfn_paths
             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)
         metafunc.parametrize("toml_name", toml_names, ids=toml_names)
 
 
@@ -45,23 +46,24 @@ def test_load_v1(dfn_name):
         (DFN_DIR / "common.dfn").open() as common_file,
         (DFN_DIR / f"{dfn_name}.dfn").open() as dfn_file,
     ):
-        common, _ = load(common_file)
+        common, _ = _load_common(common_file)
         dfn = load(dfn_file, name=dfn_name, common=common)
-        assert any(dfn)
+        assert any(dfn.fields)
 
 
 @requires_pkg("boltons")
 def test_load_v2(toml_name):
     with (TOML_DIR / f"{toml_name}.toml").open(mode="rb") as toml_file:
-        toml = Dfn.load(toml_file, name=toml_name, version=2)
-        assert any(toml)
+        dfn = load(toml_file, name=toml_name, format="toml")
+        assert any(dfn.fields)
 
 
 @requires_pkg("boltons")
-@pytest.mark.parametrize("version", list(VERSIONS.keys()))
-def test_load_all(version):
-    dfns = Dfn.load_all(VERSIONS[version], version=version)
-    assert any(dfns)
+@pytest.mark.parametrize("schema_version", list(SPEC_DIRS.keys()))
+def test_load_all(schema_version):
+    path = SPEC_DIRS[schema_version]
+    dfns = load_all(path)
+    assert all(any(dfn.fields) for dfn in dfns.values())
 
 
 @requires_pkg("boltons")
@@ -90,37 +92,34 @@ def test_load_tree():
         assert gwf_data["name"] == "gwf"
         assert gwf_data["parent"] == "sim"
 
-        # Test hierarchy enforcement and completeness
-        dfns = Dfn.load_all(tmp_path, version=2)
-        roots = [name for name, dfn in dfns.items() if not dfn.get("parent")]
-        assert len(roots) == 1
-        assert roots[0] == "sim"
-
+        dfns = load_all(tmp_path)
+        root = load_tree(tmp_path)
+        roots = []
         for dfn in dfns.values():
-            parent = dfn.get("parent")
-            if parent:
-                assert parent in dfns
-
-        # Test tree building and navigation
-        tree = Dfn.load_tree(tmp_path, version=2)
-        assert "sim" in tree
-        assert tree["sim"]["name"] == "sim"
-
-        for model_type in ["gwf", "gwt", "gwe"]:
-            if model_type in tree["sim"]:
-                assert tree["sim"][model_type]["name"] == model_type
-                assert tree["sim"][model_type]["parent"] == "sim"
-
-        if "gwf" in tree["sim"]:
+            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 tree["sim"]["gwf"].keys()
-                if k.startswith("gwf-") and isinstance(tree["sim"]["gwf"][k], dict)
+                k for k in pkgs if k.startswith("gwf-") and isinstance(pkgs[k], dict)
             ]
             assert len(gwf_packages) > 0
 
-            if "gwf-dis" in tree["sim"]["gwf"]:
-                dis = tree["sim"]["gwf"]["gwf-dis"]
-                assert dis["name"] == "gwf-dis"
-                assert dis["parent"] == "gwf"
-                assert "options" in dis or "dimensions" in dis
+            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

@@ -257,7 +257,7 @@ def _record_fields() -> Fields:
                     "block": block,
                     "description": description,
                     "default": default,
-                    **_field,
+                    **field_dict,
                 }
             )
 
@@ -337,55 +337,52 @@ def map(
     raise ValueError(f"Unsupported schema version: {schema_version}. Expected 1 or 2.")
 
 
-def load(f: str | PathLike, **kwargs) -> Dfn:
+def load(f, name: str, format: str = "dfn", **kwargs) -> Dfn:
     """Load a MODFLOW 6 definition file."""
-    path = Path(f).expanduser().resolve()
-    if path.suffix == ".dfn":
-        fields, meta = parse_dfn(path, **kwargs)
-        blocks = (
-            {}
-            if "common" in path.stem
-            else {
-                block_name: {field["name"]: FieldV1.from_dict(field) for field in block}
-                for block_name, block in groupby(
-                    fields.values(), lambda field: field["block"]
-                )
-            }
-        )
+    if format == "dfn":
+        fields, meta = parse_dfn(f, **kwargs)
+        blocks = {
+            block_name: {field["name"]: FieldV1.from_dict(field) for field in block}
+            for block_name, block in groupby(
+                fields.values(), lambda field: field["block"]
+            )
+        }
         return Dfn(
-            name=path.stem,
+            name=name,
             schema_version=Version("1"),
             parent=try_parse_parent(meta),
             advanced=is_advanced_package(meta),
             multi=is_multi_package(meta),
             blocks=blocks,
         )
-    elif path.suffix == ".toml":
-        with path.open("rb") as file:
-            return Dfn(**tomli.load(file))
-    raise ValueError(f"Unsupported file type: {path.suffix}. Expected .dfn or .toml.")
+    elif format == "toml":
+        return Dfn(name=name, **tomli.load(f))
+    raise ValueError(f"Unsupported format: {format}. Expected 'dfn' or 'toml'.")
 
 
-def _load_common(path: str | PathLike) -> Fields:
-    common_path = Path(path).expanduser().resolve()
-    if not common_path.is_file():
-        return {}
-    common, _ = parse_dfn(common_path)
+def _load_common(f) -> Fields:
+    common, _ = parse_dfn(f)
     return common
 
 
-def load_all(dfndir: str | PathLike) -> Dfns:
+def load_all(path: str | PathLike) -> Dfns:
     """Load a MODFLOW 6 specification from definition files in a directory."""
     exclude = ["common", "flopy"]
-    dfndir = Path(dfndir).expanduser().resolve()
-    dfns = {p.stem: p for p in dfndir.glob("*.dfn") if p.stem not in exclude}
-    tomls = {p.stem: p for p in dfndir.glob("*.toml") if p.stem not in exclude}
-    if dfns:
-        common = _load_common(dfndir / "common.dfn")
-        return {path.stem: load(path, common=common) for path in dfns.values()}
-    if tomls:
-        return {path.stem: load(path) for path in tomls.values()}
-    return {}  # TODO: raise instead?
+    path = Path(path).expanduser().resolve()
+    dfn_paths = {p.stem: p for p in path.glob("*.dfn") if p.stem not in exclude}
+    toml_paths = {p.stem: p for p in path.glob("*.toml") if p.stem not in exclude}
+    dfns = {}
+    if dfn_paths:
+        with (path / "common.dfn").open() as f:
+            common = _load_common(f)
+        for dfn_name, dfn_path in dfn_paths.items():
+            with dfn_path.open() as f:
+                dfns[dfn_name] = load(f, name=dfn_name, common=common, format="dfn")
+    if toml_paths:
+        for toml_name, toml_path in toml_paths.items():
+            with toml_path.open() as f:
+                dfns[toml_name] = load(f, name=toml_name, format="toml")
+    return dfns
 
 
 def infer_tree(dfns: Dfns) -> Dfn:

modflow_devtools/dfn/parse.py

@@ -1,6 +1,4 @@
 from ast import literal_eval
-from os import PathLike
-from pathlib import Path
 from typing import Any
 from warnings import warn
 
@@ -83,17 +81,14 @@ def is_multi_package(meta: list[str]) -> bool:
     return any("multi-package" in m for m in meta)
 
 
-def parse_dfn(
-    path: str | PathLike,
-    common: dict | None = None,
-) -> tuple[OMD, list[str]]:
+def parse_dfn(f, common: dict | None = None) -> tuple[OMD, list[str]]:
     """
     Parse a DFN file into an ordered dict of fields and a list of metadata.
 
     Parameters
     ----------
-    path : str | PathLike
-        Path to the DFN file to parse.
+    f : readable file-like
+        A file-like object to read the DFN file from.
     common : dict, optional
         A dictionary of common variable definitions to use for
         description substitutions, by default None.
@@ -123,7 +118,7 @@ def parse_dfn(
     fields: list = []
     metadata: list = []
 
-    for line in Path(path).expanduser().resolve().open():
+    for line in f:
         # parse metadata line
         if (line := line.strip()).startswith("#"):
             _, sep, tail = line.partition("flopy")

modflow_devtools/dfn/schema/v1.py

@@ -19,4 +19,5 @@ class FieldV1(Field):
     @classmethod
     def from_dict(cls, d: dict) -> "FieldV1":
         """Create a FieldV1 instance from a dictionary."""
-        return cls(**{k: v for k, v in d.items() if k in cls.__annotations__.keys()})
+        keys = list(cls.__annotations__.keys()) + list(Field.__annotations__.keys())
+        return cls(**{k: v for k, v in d.items() if k in keys})

modflow_devtools/dfn/schema/v2.py

@@ -15,4 +15,5 @@ class FieldV2(Field):
     @classmethod
     def from_dict(cls, d: dict) -> "FieldV2":
         """Create a FieldV2 instance from a dictionary."""
-        return cls(**{k: v for k, v in d.items() if k in cls.__annotations__.keys()})
+        keys = list(cls.__annotations__.keys()) + list(Field.__annotations__.keys())
+        return cls(**{k: v for k, v in d.items() if k in keys})

modflow_devtools/dfn2toml.py

@@ -6,8 +6,10 @@
 from pathlib import Path
 
 import tomli_w as tomli
+from boltons.iterutils import remap
 
 from modflow_devtools.dfn import load_all, map
+from modflow_devtools.misc import drop_none_or_empty
 
 # mypy: ignore-errors
 
@@ -22,7 +24,11 @@ def convert(indir: PathLike, outdir: PathLike, schema_version: str = "2") -> Non
     }
     for dfn_name, dfn in dfns.items():
         with Path.open(outdir / f"{dfn_name}.toml", "wb") as f:
-            tomli.dump(asdict(dfn), f)
+            dfn_dict = asdict(dfn)
+            # TODO if we start using c/attrs, swap
+            # this for a custom unstructuring hook
+            dfn_dict["schema_version"] = str(dfn_dict["schema_version"])
+            tomli.dump(remap(dfn_dict, visit=drop_none_or_empty), f)
 
 
 if __name__ == "__main__":

modflow_devtools/misc.py

@@ -581,3 +581,13 @@ def try_literal_eval(value: str) -> Any:
         return literal_eval(value)
     except (SyntaxError, ValueError):
         return value
+
+
+def drop_none_or_empty(path, key, value):
+    """
+    Drop dictionary items with None or empty string values.
+    For use with `boltons.iterutils.remap`.
+    """
+    if value is None or value == "":
+        return False
+    return True

modflow_devtools/models.py

@@ -21,13 +21,7 @@
 from pooch import Pooch
 
 import modflow_devtools
-from modflow_devtools.misc import get_model_paths
-
-
-def _drop_none_or_empty(path, key, value):
-    if value is None or value == "":
-        return False
-    return True
+from modflow_devtools.misc import drop_none_or_empty, get_model_paths
 
 
 def _model_sort_key(k) -> int:
@@ -421,7 +415,7 @@ def index(
 
         with self._registry_file_path.open("ab+") as registry_file:
             tomli_w.dump(
-                remap(dict(sorted(files.items())), visit=_drop_none_or_empty),
+                remap(dict(sorted(files.items())), visit=drop_none_or_empty),
                 registry_file,
             )