feat(models): add example scenario mapping (#198)

Expose the grouping of example models into scenarios

Author: wpbonelli

autotest/test_models.py

@@ -35,6 +35,13 @@ def test_models(model_name, files):
         assert any(Path(f).name == "mfsim.nam" for f in files)
 
 
+@pytest.mark.parametrize("example_name, model_names", models.get_examples().items())
+def test_get_examples(example_name, model_names):
+    assert example_name in models.EXAMPLES
+    for model_name in model_names:
+        assert model_name in models.MODELMAP
+
+
 @pytest.mark.parametrize("model_name, files", list(islice(MODELMAP.items(), TAKE)))
 def test_copy_to(model_name, files, tmp_path):
     workspace = models.copy_to(tmp_path, model_name)

modflow_devtools/make_registry.py

@@ -9,7 +9,10 @@
 from modflow_devtools.misc import get_model_paths
 from modflow_devtools.models import BASE_URL
 
-REGISTRY_PATH = Path(__file__).parent / "registry" / "registry.toml"
+REGISTRY_DIR = Path(__file__).parent / "registry"
+REGISTRY_PATH = REGISTRY_DIR / "registry.toml"
+MODELMAP_PATH = REGISTRY_DIR / "models.toml"
+EXAMPLES_PATH = REGISTRY_DIR / "examples.toml"
 
 
 def _sha256(path: Path) -> str:
@@ -28,21 +31,16 @@ def _sha256(path: Path) -> str:
 
 def write_registry(
     path: str | PathLike,
-    registry_path: str | PathLike,
     url: str,
     append: bool = False,
 ):
     path = Path(path).expanduser().absolute()
-    registry_path = Path(registry_path).expanduser().absolute()
-    modelmap_path = registry_path.parent / "models.toml"
-
     if not path.is_dir():
         raise NotADirectoryError(f"Path {path} is not a directory.")
-    if not registry_path.exists():
-        registry_path.parent.mkdir(parents=True, exist_ok=True)
 
     registry: dict[str, dict[str, str | None]] = {}
     modelmap: dict[str, list[str]] = {}
+    examples: dict[str, list[str]] = {}
     exclude = [".DS_Store", "compare"]
     if is_zip := url.endswith((".zip", ".tar")):
         registry[url.rpartition("/")[2]] = {"hash": None, "url": url}
@@ -60,10 +58,13 @@ def _find_examples_dir(p):
         # then the model names could correspond directly to directory names.
         model_path = model_path.expanduser().absolute()
         base_path = _find_examples_dir(model_path) if is_zip else path
-        model_name = (
-            str(model_path.relative_to(base_path)).replace("/", "_").replace("-", "_")
-        )
+        rel_path = model_path.relative_to(base_path)
+        model_name = str(rel_path).replace("/", "_").replace("-", "_")
         modelmap[model_name] = []
+        if is_zip:
+            if rel_path.parts[0] not in examples:
+                examples[rel_path.parts[0]] = []
+            examples[rel_path.parts[0]].append(model_name)
         for p in model_path.glob("*"):
             if not p.is_file() or any(e in p.name for e in exclude):
                 continue
@@ -85,15 +86,19 @@ def drop_none_or_empty(path, key, value):
             return False
         return True
 
-    with registry_path.open("ab+" if append else "wb") as registry_file:
+    REGISTRY_DIR.mkdir(parents=True, exist_ok=True)
+    with REGISTRY_PATH.open("ab+" if append else "wb") as registry_file:
         tomli.dump(
             remap(dict(sorted(registry.items())), visit=drop_none_or_empty),
             registry_file,
         )
 
-    with modelmap_path.open("ab+" if append else "wb") as modelmap_file:
+    with MODELMAP_PATH.open("ab+" if append else "wb") as modelmap_file:
         tomli.dump(dict(sorted(modelmap.items())), modelmap_file)
 
+    with EXAMPLES_PATH.open("ab+" if append else "wb") as examples_file:
+        tomli.dump(dict(sorted(examples.items())), examples_file)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Make a registry of example models.")
@@ -113,4 +118,4 @@ def drop_none_or_empty(path, key, value):
     args = parser.parse_args()
     path = Path(args.path)
     url = args.url if args.url else BASE_URL
-    write_registry(path=path, registry_path=REGISTRY_PATH, url=url, append=args.append)
+    write_registry(path=path, url=url, append=args.append)

modflow_devtools/models.py

@@ -18,6 +18,7 @@
 REGISTRY_ANCHOR = f"{modflow_devtools.__name__}.registry"
 REGISTRY_FILE_NAME = "registry.toml"
 MODELMAP_FILE_NAME = "models.toml"
+EXAMPLES_FILE_NAME = "examples.toml"
 
 # the mf6 examples release is our base url
 BASE_URL = "https://github.com/MODFLOW-ORG/modflow6-examples/releases/download/current"
@@ -26,7 +27,8 @@
 # set up the pooch
 FETCHERS = {}
 REGISTRY: dict[str, str] = {}
-MODELMAP: dict[str, list[Path]] = {}
+MODELMAP: dict[str, list[str]] = {}
+EXAMPLES: dict[str, list[str]] = {}
 POOCH = pooch.create(
     path=pooch.os_cache(modflow_devtools.__name__.replace("_", "-")),
     base_url=BASE_URL,
@@ -86,6 +88,22 @@ def _fetch_zip(zip_name):
     )
 
 
+try:
+    with pkg_resources.open_binary(
+        REGISTRY_ANCHOR, EXAMPLES_FILE_NAME
+    ) as examples_file:
+        EXAMPLES = tomli.load(examples_file)
+except:  # noqa: E722
+    warn(
+        f"No examples file '{EXAMPLES_FILE_NAME}' "
+        f"in module '{REGISTRY_ANCHOR}' resources"
+    )
+
+
+def get_examples() -> dict[str, list[str]]:
+    return EXAMPLES
+
+
 def get_registry() -> dict[str, str]:
     return POOCH.registry