refactor(models): introduce model registry classes (#210)

Introduce a ModelRegistry base class with LocalRegistry and PoochRegistry variants. Addresses an open topic from #134:

    One question is whether the last two use cases can be merged somehow, i.e. how to give the models API the ability to access local models. It is often necessary to modify a local model during development, and it is inconvenient to have to switch between the models API and the fixtures for this.

The aim here is for an mf6 developer to be able to transparently swap in a local registry in order to access local models, where otherwise the "official" registry would be the (user- and developer-facing) default.

This changeset should not break any existing usages of the preexisting model registry, fingers crossed.

Author: wpbonelli

autotest/test_models.py

@@ -9,25 +9,23 @@
 
 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())
-)
+@pytest.mark.parametrize("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 +37,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,21 @@
 # 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 `ModelRegistry`. There is one "official" `PoochRegistry`, aimed at users and developers — developers may create `LocalRegistry` instances to load models from the local filesystem.
+
+This module leans heavily on [Pooch](https://www.fatiando.org/pooch/latest/index.html), but it is an independent layer on top with strong opinions about how to train (configure) the fetch-happy friend.
+
+## `ModelRegistry`
+
+Registries expose the following properties:
+
+- `path`: the data path
+- `files`: a map of files to file info
+- `models`: a map of models to files
+- `examples`: a map of example scenarios to models
+
+An *example* is a set of models which run in a particular order.
+
+The default `PoochRegistry` is available at `modflow_devtools.models.DEFAULT_REGISTRY`. Its `path` is the pooch cache. Values in the `files` are dictionaries including a hash and url. Configuring the default registry is a developer task — see the instructions on [creating a registry](#creating-a-registry) below.
 
 ## Listing models
 
@@ -51,7 +66,19 @@ with TemporaryDirectory() as td:
 
 If the target directory doesn't exist, it will be created.
 
-## Developers
+## Creating a registry
+
+### Local registries
+
+A `LocalRegistry` accepts a `path` on initialization. This must be a directory containing model subdirectories at arbitrary depth. Model subdirectories are identified by the presence of a namefile matching `namefile_pattern`. By default `namefile_pattern="mfsim.nam"`, causing only MODFLOW 6 models to be returned.
+
+For instance, to load all MODFLOW models (pre-MF6 as well):
+
+```python
+registry = LocalRegistry("path/to/models", namefile_pattern="*.nam")
+```
+
+### Pooch 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,
     )