refactor(models): introduce model registry class

Author: wpbonelli

autotest/test_models.py autotest/test_registry_default.pyautotest/test_models.py renamed to autotest/test_registry_default.py

@@ -9,25 +9,25 @@
 
 TAKE = 5 if is_in_ci() else None
 PROJ_ROOT = Path(__file__).parents[1]
-MODELMAP_PATH = PROJ_ROOT / "modflow_devtools" / "registry" / models.MODELMAP_FILE_NAME
-with MODELMAP_PATH.open("rb") as f:
-    MODELMAP = tomli.load(f)
+MODELS_PATH = PROJ_ROOT / "modflow_devtools" / "registry" / "models.toml"
+MODELS = tomli.load(MODELS_PATH.open("rb"))
+REGISTRY = models._DEFAULT_REGISTRY
 
 
-def test_registry():
-    registry = models.get_registry()
-    assert registry is not None, "Registry was not loaded"
-    assert registry is models.POOCH.registry
-    assert any(registry), "Registry is empty"
+def test_files():
+    files = models.get_files()
+    assert files is not None, "Files not loaded"
+    assert files is REGISTRY.files
+    assert any(files), "Registry is empty"
 
 
 @pytest.mark.parametrize(
-    "model_name, files", MODELMAP.items(), ids=list(MODELMAP.keys())
+    "model_name, files", MODELS.items(), ids=list(MODELS.keys())
 )
 def test_models(model_name, files):
     model_names = list(models.get_models().keys())
     assert model_name in model_names, f"Model {model_name} not found in model map"
-    assert files == models.MODELMAP[model_name], (
+    assert files == REGISTRY.models[model_name], (
         f"Files for model {model_name} do not match"
     )
     if "mf6" in model_name:
@@ -39,16 +39,16 @@ def test_models(model_name, files):
     models.get_examples().items(),
     ids=list(models.get_examples().keys()),
 )
-def test_get_examples(example_name, model_names):
-    assert example_name in models.EXAMPLES
+def test_examples(example_name, model_names):
+    assert example_name in models.get_examples()
     for model_name in model_names:
-        assert model_name in models.MODELMAP
+        assert model_name in REGISTRY.models
 
 
 @pytest.mark.parametrize(
     "model_name, files",
-    list(islice(MODELMAP.items(), TAKE)),
-    ids=list(MODELMAP.keys())[:TAKE],
+    list(islice(MODELS.items(), TAKE)),
+    ids=list(MODELS.keys())[:TAKE],
 )
 def test_copy_to(model_name, files, tmp_path):
     workspace = models.copy_to(tmp_path, model_name, verbose=True)

docs/md/models.md

@@ -1,6 +1,8 @@
 # Models API
 
-The `modflow_devtools.models` module provides programmatic access to MODFLOW 6 example models via [Pooch](https://www.fatiando.org/pooch/latest/index.html).
+The `modflow_devtools.models` module provides programmatic access to MODFLOW 6 example models via a `Models` registry class. There is one "official" registry, aimed at users and developers — developers may create additional instances e.g. to load models from the local filesystem.
+
+This module leans heavily on [Pooch](https://www.fatiando.org/pooch/latest/index.html), but it has strong opinions about how to train (configure) it.
 
 ## Listing models
 
@@ -51,7 +53,7 @@ with TemporaryDirectory() as td:
 
 If the target directory doesn't exist, it will be created.
 
-## Developers
+## Creating a registry
 
 The `make_registry.py` script is responsible for generating a registry text file and a mapping between files and models. This script should be run in the CI pipeline at release time before the package is built. The generated registry file and model mapping are used to create a pooch instance for fetching model files, and should be distributed with the package.
 

modflow_devtools/make_registry.py

@@ -1,124 +1,6 @@
 import argparse
-import hashlib
-from os import PathLike
-from pathlib import Path
-
-import tomli_w as tomli
-from boltons.iterutils import remap
-
-from modflow_devtools.misc import get_model_paths
-from modflow_devtools.models import BASE_URL
-
-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:
-    """
-    Compute the SHA256 hash of the given file.
-    Reference: https://stackoverflow.com/a/44873382/6514033
-    """
-    h = hashlib.sha256()
-    b = bytearray(128 * 1024)
-    mv = memoryview(b)
-    with path.open("rb", buffering=0) as f:
-        for n in iter(lambda: f.readinto(mv), 0):
-            h.update(mv[:n])
-    return h.hexdigest()
-
-
-def write_registry(
-    path: str | PathLike,
-    url: str,
-    prefix: str = "",
-    append: bool = False,
-    namefile: str = "mfsim.nam",
-):
-    """
-    Make registry files for a directory of models.
-
-    The directory may contain model subdirectories
-    at arbitrary depth. Model input subdirectories
-    are identified by the presence of a namefile
-    matching the provided pattern. A prefix may be
-    specified for model names to avoid collisions.
-    The registry files are written to the registry
-    folder alongside this script. Typically, this
-    function will run once or more in append mode
-    to iteratively create a registry.
-
-    Parameters
-    ----------
-    path : str | PathLike
-        Path to the directory containing the models.
-    url : str
-        Base URL for the models.
-    prefix : str
-        Prefix to add to model names.
-    append : bool
-        Append to the registry files instead of overwriting them.
-    namefile : str
-        Namefile pattern to look for in the model directories.
-    """
-    path = Path(path).expanduser().absolute()
-    if not path.is_dir():
-        raise NotADirectoryError(f"Path {path} is not a directory.")
-
-    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}
-
-    model_paths = get_model_paths(path, namefile=namefile)
-    for model_path in model_paths:
-        model_path = model_path.expanduser().absolute()
-        rel_path = model_path.relative_to(path)
-        parts = [prefix, *list(rel_path.parts)] if prefix else list(rel_path.parts)
-        model_name = "/".join(parts)
-        modelmap[model_name] = []
-        if is_zip:
-            name = rel_path.parts[0]
-            if name not in examples:
-                examples[name] = []
-            examples[name].append(model_name)
-        for p in model_path.rglob("*"):
-            if not p.is_file() or any(e in p.name for e in exclude):
-                continue
-            if is_zip:
-                relpath = p.expanduser().absolute().relative_to(path)
-                name = "/".join(relpath.parts)
-                url_ = url
-                hash = None
-            else:
-                relpath = p.expanduser().absolute().relative_to(path)
-                name = "/".join(relpath.parts)
-                url_ = f"{url}/{relpath!s}"
-                hash = _sha256(p)
-            registry[name] = {"hash": hash, "url": url_}
-            modelmap[model_name].append(name)
-
-    def drop_none_or_empty(path, key, value):
-        if value is None or value == "":
-            return False
-        return True
-
-    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:
-        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)
 
+import modflow_devtools.models as models
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Make a registry of models.")
@@ -137,7 +19,7 @@ def drop_none_or_empty(path, key, value):
         "-u",
         type=str,
         help="Base URL for models.",
-        default=BASE_URL,
+        default=models._DEFAULT_BASE_URL,
     )
     parser.add_argument(
         "--namefile",
@@ -147,10 +29,13 @@ def drop_none_or_empty(path, key, value):
         default="mfsim.nam",
     )
     args = parser.parse_args()
-    write_registry(
+    if not args.append:
+        models._DEFAULT_REGISTRY = models.PoochRegistry(
+            base_url=args.url, env=models._DEFAULT_ENV
+        )
+    models._DEFAULT_REGISTRY.index(
         path=args.path,
         url=args.url,
         prefix=args.prefix,
-        append=args.append,
         namefile=args.namefile,
     )