| """Render Jinja2 templates with a BundleSpec. |
| |
| Templates are shipped as package data under `casmd/bundle/templates/`. |
| """ |
| from __future__ import annotations |
| from dataclasses import asdict |
| from importlib import resources |
| from typing import Any |
|
|
| from jinja2 import Environment, BaseLoader, TemplateNotFound |
|
|
| from casmd.bundle.spec import BundleSpec |
|
|
|
|
| _TEMPLATE_NAMES = ( |
| "step1_minimization.mdp", |
| "step2_nvt.mdp", |
| "step3_npt.mdp", |
| "step4_production.mdp", |
| "submit.sh", |
| "run_md.sh", |
| "analyze.py", |
| "README.md", |
| ) |
|
|
|
|
| class _PackageLoader(BaseLoader): |
| """Load Jinja2 templates from casmd.bundle.templates package data.""" |
|
|
| def get_source(self, environment, template: str): |
| try: |
| res = resources.files("casmd.bundle.templates").joinpath(f"{template}.j2") |
| text = res.read_text(encoding="utf-8") |
| except (FileNotFoundError, ModuleNotFoundError): |
| raise TemplateNotFound(template) |
| return text, None, lambda: True |
|
|
|
|
| _ENV = Environment(loader=_PackageLoader(), keep_trailing_newline=True) |
|
|
|
|
| def list_template_names() -> list[str]: |
| return list(_TEMPLATE_NAMES) |
|
|
|
|
| def render(name: str, spec: BundleSpec, *, extra: dict[str, Any] | None = None) -> str: |
| """Render template `name` using a BundleSpec (plus optional extra context).""" |
| if name not in _TEMPLATE_NAMES: |
| raise KeyError(f"unknown template {name!r}; known: {_TEMPLATE_NAMES}") |
| ctx: dict[str, Any] = asdict(spec) |
| |
| ctx["production_steps"] = spec.production_steps |
| ctx["nvt_steps"] = spec.nvt_steps |
| ctx["npt_steps"] = spec.npt_steps |
| ctx["output_every_steps"] = spec.output_every_steps |
| if extra: |
| ctx.update(extra) |
| template = _ENV.get_template(name) |
| return template.render(**ctx) |
|
|
