WIP dvc: implement vars templating in multistage dvc file

dvc/cli.py

@@ -33,6 +33,7 @@
     plots,
     remote,
     remove,
+    render,
     repro,
     root,
     run,
@@ -81,6 +82,7 @@
     plots,
     experiments,
     check_ignore,
+    render,
 ]
 
 

dvc/command/render.py

@@ -0,0 +1,62 @@
+import argparse
+import logging
+
+from dvc.command.base import CmdBase, append_doc_link
+from dvc.utils.serialize._yaml import dumps_yaml, render_dvc_template
+
+logger = logging.getLogger(__name__)
+
+
+class CmdRender(CmdBase):
+    def run(self):
+        with open(self.args.path, encoding="utf-8") as fd:
+            text = fd.read()
+        dvc_dict, vars_dict = render_dvc_template(text)
+        vars_dict = {"vars": vars_dict}
+        vars_str = dumps_yaml(vars_dict)
+
+        if self.args.only_vars:
+            logger.info(vars_str)
+            return 0
+
+        if self.args.stage is not None:
+            dvc_dict = dvc_dict["stages"][self.args.stage]
+        dvc_str = dumps_yaml({"stages": {self.args.stage: dvc_dict}})
+
+        if self.args.only_stages:
+            logger.info(dvc_str)
+            return 0
+
+        logger.info("\n".join([vars_str, dvc_str]))
+        return 0
+
+
+def add_parser(subparsers, parent_parser):
+    RENDER_HELP = "Render templated dvc.yaml"
+    render_parser = subparsers.add_parser(
+        "render",
+        parents=[parent_parser],
+        description=append_doc_link(RENDER_HELP, "render"),
+        help=RENDER_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    render_parser.add_argument(
+        "--path", default="dvc.yaml", help="Path to dvc.yaml file",
+    )
+    render_parser.add_argument(
+        "--stage", "-s", default=None, help="Only render a specified stage",
+    )
+    render_parser.add_argument(
+        "--only-vars",
+        action="store_true",
+        default=False,
+        help="Only render the `vars` component of the dvc.yaml",
+    )
+    render_parser.add_argument(
+        "--only-stages",
+        action="store_true",
+        default=False,
+        help="Only render the `stages` component of the dvc.yaml",
+    )
+
+    render_parser.set_defaults(func=CmdRender)

dvc/utils/serialize/_yaml.py

@@ -3,6 +3,7 @@
 from contextlib import contextmanager
 
 from funcy import reraise
+from jinja2 import Template
 from ruamel.yaml import YAML
 from ruamel.yaml.error import YAMLError
 
@@ -14,14 +15,77 @@ def __init__(self, path):
         super().__init__(path, "YAML file structure is corrupted")
 
 
-def load_yaml(path, tree=None):
-    return _load_data(path, parser=parse_yaml, tree=tree)
+def recursive_render(tpl, values, max_passes=100):
+    """This is a bit of black magic to recursivly render a
+    template. Adaped from:
+
+      https://stackoverflow.com/questions/8862731/jinja-nested-rendering-on-variable-content
+
+    Args:
+      tpl: Template string
+      values: dict of values. Importantly this dict can contain
+          values that are themselves {{ placeholders }}
+      max_passes: Limits the number of times we loop over the
+          template.
+
+    Returns:
+      rendered template.
+    """
+    prev = tpl
+    for _ in range(max_passes):
+        curr = Template(prev).render(**values)
+        if curr != prev:
+            prev = curr
+        else:
+            return curr
+    raise RecursionError("Max resursion depth reached")
+
+
+def render_vars(dvc_dict):
+    vars_dict = dvc_dict["vars"]
+    vars_template = dumps_yaml(vars_dict)
+
+    rendered_vars = recursive_render(vars_template, vars_dict)
+    return rendered_vars
+
+
+def render_dvc(dvc_dict, vars_dict):
+    dvc_template = dumps_yaml(dvc_dict)
+    rendered_dvc = Template(dvc_template).render(**vars_dict)
+    return rendered_dvc
+
+
+def render_dvc_template(text):
+
+    yaml = YAML(typ="safe")
+    dvc_dict = yaml.load(text) or {}
+    if "vars" in dvc_dict:
+        vars_dict = yaml.load(render_vars(dvc_dict))
+
+        del dvc_dict["vars"]
+        rendered_dvc = yaml.load(render_dvc(dvc_dict, vars_dict)) or {}
+
+    return rendered_dvc, vars_dict
 
 
 def parse_yaml(text, path, typ="safe"):
     yaml = YAML(typ=typ)
     with reraise(YAMLError, YAMLFileCorruptedError(path)):
-        return yaml.load(text) or {}
+        result = yaml.load(text) or {}
+
+    if "vars" in result:
+        try:
+            result, _ = render_dvc_template(
+                text
+            )  # yaml.load(text, Loader=SafeLoader) or {}
+        except Exception as exc:
+            raise YAMLFileCorruptedError(path) from exc
+
+    return result
+
+
+def load_yaml(path, tree=None):
+    return _load_data(path, parser=parse_yaml, tree=tree)
 
 
 def parse_yaml_for_update(text, path):