feat: Nested Templates

.pre-commit-config.yaml

@@ -62,7 +62,7 @@ repos:
         files: src
         args: []
         additional_dependencies:
-          ["types-Jinja2", "types-toml", "typer", "importlib_resources"]
+          ["types-Jinja2", "typer", "importlib_resources"]
 
   - repo: https://github.com/codespell-project/codespell
     rev: v2.2.4

pyproject.toml

@@ -37,7 +37,8 @@ classifiers = [
 dependencies = [
     "typing_extensions >=3.7; python_version<'3.11'",
     "typer",
-    "toml",
+    "tomli; python_version<'3.11'",
+    "tomli-w",
     "jinja2",
     "in_place",
     "importlib_resources>=1.3.0; python_version<'3.9'",  # for resources in schema
@@ -109,7 +110,7 @@ dependencies = [
 ]
 
 [tool.hatch.envs.dev.scripts]
-test = "pytest -ra"
+test = "pytest -ra {args}"
 
 [[tool.hatch.envs.dev.matrix]]
 python = ["3.7", "3.8", "3.9", "3.10", "3.11", "pypy3.8"]
@@ -180,6 +181,14 @@ warn_unreachable = true
 module = 'in_place'
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = 'tomli'
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = 'tomli_w'
+ignore_missing_imports = true
+
 [tool.ruff]
 select = [
   "E", "F", "W", # flake8

src/mapyde/utils.py

@@ -11,7 +11,12 @@
 import unicodedata
 from pathlib import Path
 
-import toml
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
+import tomli_w
 from jinja2 import Environment, FileSystemLoader, Template
 
 from mapyde import prefix
@@ -51,13 +56,13 @@ def render_string(blob: str, variables: ImmutableConfig | None = None) -> str:
     variables = variables or {}
     tpl = Template(blob)
     return tpl.render(
-        PWD=os.getenv("PWD"),
-        USER=os.getenv("USER"),
-        MAPYDE_DATA=prefix.data,  # type: ignore[attr-defined]
-        MAPYDE_CARDS=prefix.cards,  # type: ignore[attr-defined]  # pylint: disable=no-member
-        MAPYDE_LIKELIHOODS=prefix.likelihoods,  # type: ignore[attr-defined]  # pylint: disable=no-member
-        MAPYDE_SCRIPTS=prefix.scripts,  # type: ignore[attr-defined]  # pylint: disable=no-member
-        MAPYDE_TEMPLATES=prefix.templates,  # type: ignore[attr-defined]  # pylint: disable=no-member
+        PWD=Path(os.getenv("PWD", ".")).as_posix(),
+        USER=os.getenv("USER", "USER"),
+        MAPYDE_DATA=prefix.data.as_posix(),  # type: ignore[attr-defined]
+        MAPYDE_CARDS=prefix.cards.as_posix(),  # type: ignore[attr-defined]  # pylint: disable=no-member
+        MAPYDE_LIKELIHOODS=prefix.likelihoods.as_posix(),  # type: ignore[attr-defined]  # pylint: disable=no-member
+        MAPYDE_SCRIPTS=prefix.scripts.as_posix(),  # type: ignore[attr-defined]  # pylint: disable=no-member
+        MAPYDE_TEMPLATES=prefix.templates.as_posix(),  # type: ignore[attr-defined]  # pylint: disable=no-member
         **variables,
     )
 
@@ -78,29 +83,46 @@ def load_config(filename: str, cwd: str = ".") -> T.Any:
 
     tpl = env.get_template(filename)
     assert tpl.filename
-    with Path(tpl.filename).open(encoding="utf-8") as fpointer:
-        return toml.load(fpointer)
+    with Path(tpl.filename).open("rb") as fpointer:
+        return tomllib.load(fpointer)
 
 
-def build_config(user: MutableConfig) -> T.Any:
+def build_config(user: MutableConfig, depth: int = 0) -> T.Any:
     """
     Function to build a configuration from a user-provided toml configuration on top of the base/template one.
-    """
 
-    template_path = Path(
-        render_string(
-            user["base"].get("template", "{{MAPYDE_TEMPLATES}}/defaults.toml")
-        )
-    )
+    The templates can be further nested (this function is recursive) up to a maximum (non-configurable) depth of 10.
 
-    with resources.as_file(template_path) as template:
-        if not template.exists():
-            msg = f"{template_path} does not exist."
-            raise OSError(msg)
-        defaults = load_config(template.name, str(template.parent))
+    """
 
-    variables = merge(defaults, user)
-    return toml.loads(render_string(toml.dumps(variables), variables))
+    template_str = user.get("base", {}).pop(
+        "template", str(prefix.templates / "defaults.toml")  # type: ignore[attr-defined]  # pylint: disable=no-member
+    )
+    parent = {}
+
+    if template_str:
+        if depth >= 10:
+            msg = 'Maximum template depth (10) exceeded. This is likely due to your base template not having `"template" = false` set.'
+            raise RuntimeError(msg)
+
+        template_path = Path(render_string(template_str))
+
+        with resources.as_file(template_path) as template:
+            if not template.exists():
+                msg = f"{template_path} does not exist."
+                raise OSError(msg)
+            parent = build_config(
+                load_config(template.name, str(template.parent)), depth=depth + 1
+            )
+
+    variables = merge(parent, user)
+
+    # only render the entire merged configuration, not the intermediate ones
+    return (
+        variables
+        if depth
+        else tomllib.loads(render_string(tomli_w.dumps(variables), variables))
+    )
 
 
 def output_path(config: ImmutableConfig) -> Path:

templates/defaults.toml

@@ -13,6 +13,7 @@ pythia_path = "{{MAPYDE_CARDS}}/pythia/"
 delphes_path = "{{MAPYDE_CARDS}}/delphes/"
 madspin_path = "{{MAPYDE_CARDS}}/madspin/"
 likelihoods_path = "{{MAPYDE_LIKELIHOODS}}"
+template = false
 
 [madgraph]
 skip = false

tests/test_cli.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import time
+
+import pytest
+from typer.testing import CliRunner
+
+import mapyde
+from mapyde.cli import app
+
+
+@pytest.fixture()
+def runner():
+    return CliRunner()
+
+
+def test_version(runner):
+    start = time.time()
+    result = runner.invoke(app, ["--version"])
+    end = time.time()
+    elapsed = end - start
+    assert result.exit_code == 0
+    assert mapyde.__version__ in result.stdout
+    # make sure it took less than a second
+    assert elapsed < 1.0

tests/test_config.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import pytest
+import tomli_w
+
+import mapyde
+import mapyde.utils
+
+
+def test_template_false():
+    config = {
+        "base": {
+            "path": "/data/users/{{USER}}/SUSY",
+            "output": "mytag",
+            "template": False,
+        },
+        "madgraph": {
+            "proc": {"name": "charginos", "card": "{{madgraph['proc']['name']}}"},
+            "masses": {"MN2": 500},
+        },
+    }
+    built = mapyde.utils.build_config(config)
+    assert built
+    assert "{{USER}}" not in built["base"]["path"]
+    assert "madspin" not in built
+
+
+def test_template_empty():
+    config = {
+        "base": {"path": "/data/users/{{USER}}/SUSY", "output": "mytag"},
+        "madgraph": {
+            "proc": {"name": "charginos", "card": "{{madgraph['proc']['name']}}"},
+            "masses": {"MN2": 500},
+        },
+    }
+    built = mapyde.utils.build_config(config)
+    assert built
+    assert "{{USER}}" not in built["base"]["path"]
+    assert "madspin" in built
+    assert "sherpa" in built
+    assert built["madspin"]["skip"]
+    assert built["base"]["engine"] == "docker"
+    assert built["madgraph"]["params"] == "Higgsino"
+    assert built["madgraph"]["paramcard"] == "Higgsino.slha"
+
+
+def test_template_nested(tmp_path):
+    template = tmp_path / "enable_madspin.toml"
+    template.write_text(tomli_w.dumps({"madspin": {"skip": False}}))
+
+    config = {
+        "base": {
+            "engine": "fakeengine",
+            "path": "/data/users/{{USER}}/SUSY",
+            "output": "mytag",
+            "template": template.as_posix(),
+            "rendered_engine": "{{base['engine']}}",
+        },
+        "madgraph": {
+            "params": "fakeparams",
+            "proc": {"name": "charginos", "card": "{{madgraph['proc']['name']}}"},
+            "masses": {"MN2": 500},
+        },
+    }
+    built = mapyde.utils.build_config(config)
+    assert built
+    assert "{{USER}}" not in built["base"]["path"]
+    assert "madspin" in built
+    assert "sherpa" in built
+    assert not built["madspin"]["skip"]
+    assert built["base"]["engine"] == "fakeengine"
+    assert built["base"]["rendered_engine"] == "fakeengine"
+    assert built["madgraph"]["params"] == "fakeparams"
+    assert (
+        built["madgraph"]["paramcard"] == "fakeparams.slha"
+    ), "The rendering of the config happened prematurely, rather than after merging templates."
+
+
+def test_template_max_nested(tmp_path):
+    templates = tmp_path / "templates"
+    templates.mkdir()
+
+    defaults = templates / "defaults.toml"
+
+    template = mapyde.utils.load_config("defaults.toml", mapyde.prefix.templates)
+    template["base"].pop("template")  # remove the template key to force recursion
+    defaults.write_text(tomli_w.dumps(template))
+
+    with mapyde.prefix(templates.parent):
+        config = {
+            "base": {
+                "path": "/data/users/{{USER}}/SUSY",
+                "output": "mytag",
+            },
+        }
+        with pytest.raises(RuntimeError, match=r"Maximum template depth .* exceeded"):
+            mapyde.utils.build_config(config)