From 6a79111047a838b2e7daee0861570ffe01c38d02 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Wed, 14 Oct 2020 05:32:05 -0700
Subject: [PATCH] Parametrization support on dvc.yaml (#4679)

* Custom parametrization support

* log generated data

* Add tests
---
 dvc/config.py                     |   2 +
 dvc/dvcfile.py                    |   8 +++
 dvc/parsing/__init__.py           |  27 ++++++++
 dvc/parsing/context.py            |  61 ++++++++++++++++++
 dvc/parsing/interpolate.py        |  61 ++++++++++++++++++
 tests/func/test_feature_flags.py  |  34 ++++++++++
 tests/unit/test_interpolate.py    | 101 ++++++++++++++++++++++++++++++
 tests/unit/test_stage_resolver.py |  39 ++++++++++++
 8 files changed, 333 insertions(+)
 create mode 100644 dvc/parsing/__init__.py
 create mode 100644 dvc/parsing/context.py
 create mode 100644 dvc/parsing/interpolate.py
 create mode 100644 tests/func/test_feature_flags.py
 create mode 100644 tests/unit/test_interpolate.py
 create mode 100644 tests/unit/test_stage_resolver.py

diff --git a/dvc/config.py b/dvc/config.py
index 372a51d1f..090d065cd 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -224,6 +224,8 @@ SCHEMA = {
         "row_limit": All(Coerce(int), Range(1)),
         "row_cleanup_quota": All(Coerce(int), Range(0, 100)),
     },
+    # section for experimental features
+    "feature": {Optional("parametrization", default=False): Bool},
 }
 COMPILED_SCHEMA = Schema(SCHEMA)
 
diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py
index 972287362..555cd71b6 100644
--- a/dvc/dvcfile.py
+++ b/dvc/dvcfile.py
@@ -3,9 +3,11 @@ import contextlib
 import logging
 import os
 
+from funcy import log_durations
 from voluptuous import MultipleInvalid
 
 from dvc.exceptions import DvcException
+from dvc.parsing import DataResolver
 from dvc.stage import serialize
 from dvc.stage.exceptions import (
     StageFileBadNameError,
@@ -226,6 +228,12 @@ class PipelineFile(FileMixin):
     @property
     def stages(self):
         data, _ = self._load()
+
+        if self.repo.config["feature"]["parametrization"]:
+            with log_durations(logger.debug, "resolving values"):
+                resolver = DataResolver(data)
+                data = resolver.resolve()
+
         lockfile_data = self._lockfile.load()
         return StageLoader(self, data.get("stages", {}), lockfile_data)
 
diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py
new file mode 100644
index 000000000..b24345cfa
--- /dev/null
+++ b/dvc/parsing/__init__.py
@@ -0,0 +1,27 @@
+import logging
+from itertools import starmap
+
+from funcy import join
+
+from .context import Context
+from .interpolate import resolve
+
+logger = logging.getLogger(__name__)
+
+STAGES = "stages"
+
+
+class DataResolver:
+    def __init__(self, d):
+        self.context = Context()
+        self.data = d
+
+    def _resolve_entry(self, name, definition):
+        stage_d = resolve(definition, self.context)
+        logger.trace("Resolved stage data for '%s': %s", name, stage_d)
+        return {name: stage_d}
+
+    def resolve(self):
+        stages = self.data.get(STAGES, {})
+        data = join(starmap(self._resolve_entry, stages.items()))
+        return {**self.data, STAGES: data}
diff --git a/dvc/parsing/context.py b/dvc/parsing/context.py
new file mode 100644
index 000000000..3fdbc2a78
--- /dev/null
+++ b/dvc/parsing/context.py
@@ -0,0 +1,61 @@
+from collections.abc import Collection, Mapping, Sequence
+
+# for testing purpose
+# FIXME: after implementing of reading of "params".
+TEST_DATA = {
+    "__test__": {
+        "dict": {"one": 1, "two": 2, "three": "three", "four": "4"},
+        "list": [1, 2, 3, 4, 3.14],
+        "set": {1, 2, 3},
+        "tuple": (1, 2),
+        "bool": True,
+        "none": None,
+        "float": 3.14,
+        "nomnom": 1000,
+    }
+}
+
+
+class Context:
+    def __init__(self, data=None):
+        self.data = data or TEST_DATA
+
+    def select(self, key):
+        return _get_value(self.data, key)
+
+
+def _get_item(data, idx):
+    if isinstance(data, Sequence):
+        idx = int(idx)
+
+    if isinstance(data, (Mapping, Sequence)):
+        return data[idx]
+
+    raise ValueError(
+        f"Cannot get item '{idx}' from data of type '{type(data).__name__}'"
+    )
+
+
+def _get_value(data, key):
+    obj_and_attrs = key.strip().split(".")
+    value = data
+    for attr in obj_and_attrs:
+        if attr == "":
+            raise ValueError("Syntax error!")
+
+        try:
+            value = _get_item(value, attr)
+        except KeyError:
+            msg = (
+                f"Could not find '{attr}' "
+                "while substituting "
+                f"'{key}'.\n"
+                f"Interpolating with: {data}"
+            )
+            raise ValueError(msg)
+
+    if not isinstance(value, str) and isinstance(value, Collection):
+        raise ValueError(
+            f"Cannot interpolate value of type '{type(value).__name__}'"
+        )
+    return value
diff --git a/dvc/parsing/interpolate.py b/dvc/parsing/interpolate.py
new file mode 100644
index 000000000..56e72f31d
--- /dev/null
+++ b/dvc/parsing/interpolate.py
@@ -0,0 +1,61 @@
+import re
+from collections.abc import Mapping
+
+from funcy import rpartial
+
+KEYCRE = re.compile(
+    r"""
+    (?<!\\)                   # escape \${} or ${{}}
+    \$                        # starts with $
+    (?:({{)|({))              # either starts with double braces or single
+    ([\w._ \\/-]*?)           # match every char, attr access through "."
+    (?(1)}})(?(2)})           # end with same kinds of braces it opened with
+""",
+    re.VERBOSE,
+)
+
+
+def _get_matches(template):
+    return list(KEYCRE.finditer(template))
+
+
+def _resolve_value(match, context):
+    _, _, inner = match.groups()
+    return context.select(inner)
+
+
+def _str_interpolate(template, matches, context):
+    index, buf = 0, ""
+    for match in matches:
+        start, end = match.span(0)
+        buf += template[index:start] + str(_resolve_value(match, context))
+        index = end
+    return buf + template[index:]
+
+
+def _resolve_str(src: str, context):
+    matches = _get_matches(src)
+    if len(matches) == 1 and src == matches[0].group(0):
+        # replace "${enabled}", if `enabled` is a boolean, with it's actual
+        # value rather than it's string counterparts.
+        return _resolve_value(matches[0], context)
+    elif matches:
+        # but not "${num} days"
+        src = _str_interpolate(src, matches, context)
+
+    # regex already backtracks and avoids any `${` starting with
+    # backslashes(`\`). We just need to replace those by `${`.
+    return src.replace(r"\${", "${")
+
+
+def resolve(src, context):
+    Seq = (list, tuple, set)
+
+    apply_value = rpartial(resolve, context)
+    if isinstance(src, Mapping):
+        return {key: apply_value(value) for key, value in src.items()}
+    elif isinstance(src, Seq):
+        return type(src)(map(apply_value, src))
+    elif isinstance(src, str):
+        return _resolve_str(src, context)
+    return src
diff --git a/tests/func/test_feature_flags.py b/tests/func/test_feature_flags.py
new file mode 100644
index 000000000..40c900413
--- /dev/null
+++ b/tests/func/test_feature_flags.py
@@ -0,0 +1,34 @@
+from dvc.dvcfile import Dvcfile
+from dvc.utils.serialize import dumps_yaml
+from tests.unit.test_stage_resolver import (
+    RESOLVED_DVC_YAML_DATA,
+    TEMPLATED_DVC_YAML_DATA,
+)
+
+
+def test_parametrization_is_not_enabled_by_default(tmp_dir, dvc, mocker):
+    assert dvc.config["feature"]["parametrization"] is False
+
+    (tmp_dir / "dvc.yaml").write_text(dumps_yaml(RESOLVED_DVC_YAML_DATA))
+    mock = mocker.patch("dvc.dvcfile.DataResolver")
+
+    stages = list(Dvcfile(dvc, "dvc.yaml").stages)
+    mock.assert_not_called()
+    assert len(stages) == 2
+
+
+def test_parametrization_flag_when_enabled(tmp_dir, dvc, mocker):
+    dvc.config["feature"]["parametrization"] = True
+
+    mock = mocker.patch(
+        "dvc.dvcfile.DataResolver.resolve", return_value=RESOLVED_DVC_YAML_DATA
+    )
+
+    dvcfile = Dvcfile(dvc, "dvc.yaml")
+    mocker.patch.object(
+        dvcfile, "_load", return_value=[TEMPLATED_DVC_YAML_DATA, None]
+    )
+
+    stages = list(dvcfile.stages)
+    mock.assert_called_once()
+    assert len(stages) == 2
diff --git a/tests/unit/test_interpolate.py b/tests/unit/test_interpolate.py
new file mode 100644
index 000000000..e43548db3
--- /dev/null
+++ b/tests/unit/test_interpolate.py
@@ -0,0 +1,101 @@
+from math import inf, pi
+
+import pytest
+
+from dvc.parsing import Context
+from dvc.parsing.interpolate import resolve
+
+
+@pytest.mark.parametrize(
+    "template, var",
+    [
+        ("${value}", "value"),
+        ("${{item}}", "item"),
+        ("${ item }", "item"),
+        ("${{ value }}", "value"),
+    ],
+)
+@pytest.mark.parametrize(
+    "data", [True, 12, pi, None, False, 0, "0", "123", "Foobar", "", inf, 3e4]
+)
+def test_resolve_primitive_values(data, template, var):
+    assert resolve(template, Context({var: data})) == data
+
+
+@pytest.mark.parametrize(
+    "template, expected",
+    [
+        (r"\${value}", "${value}"),
+        (r"\${{value}}", "${{value}}"),
+        (r"\${ value }", "${ value }"),
+        (r"\${{ value }}", "${{ value }}"),
+        (r"\${{ value }} days", "${{ value }} days"),
+        (r"\${ value } days", "${ value } days"),
+        (r"Month of \${value}", "Month of ${value}"),
+        (r"May the \${value} be with you", "May the ${value} be with you"),
+        (
+            r"Great shot kid, that was \${value} in a ${value}",
+            "Great shot kid, that was ${value} in a value",
+        ),
+    ],
+)
+def test_escape(template, expected, mocker):
+    assert resolve(template, Context({"value": "value"})) == expected
+
+
+def test_resolve_str():
+    template = "My name is ${last}, ${first} ${last}"
+    expected = "My name is Bond, James Bond"
+    assert (
+        resolve(template, Context({"first": "James", "last": "Bond"}))
+        == expected
+    )
+
+
+def test_resolve_primitives_dict_access():
+    data = {
+        "dict": {
+            "num": 5,
+            "string": "foo",
+            "nested": {"float": pi, "string": "bar"},
+        }
+    }
+    context = Context(data)
+
+    assert resolve("${dict.num}", context) == 5
+    assert resolve("${dict.string}", context) == "foo"
+    assert resolve("${dict.nested.float}", context) == pi
+    assert resolve("${dict.nested.string}", context) == "bar"
+
+    assert resolve("Number ${dict.num}", context) == "Number 5"
+
+
+def test_resolve_primitives_list_access():
+    context = Context(
+        {
+            "dict": [
+                {"f": "f"},
+                {"fo": "fo"},
+                {"foo": "foo"},
+                {"foo": ["f", "o", "o"]},
+            ]
+        }
+    )
+
+    assert resolve("${dict.0.f}", context) == "f"
+    assert resolve("${dict.1.fo}", context) == "fo"
+    assert resolve("${dict.2.foo}", context) == "foo"
+    assert resolve("${dict.3.foo.0}", context) == "f"
+
+    assert resolve("${ dict.1.fo}${dict.3.foo.1}bar", context) == "foobar"
+
+
+def test_resolve_collection():
+    from .test_stage_resolver import (
+        CONTEXT_DATA,
+        RESOLVED_DVC_YAML_DATA,
+        TEMPLATED_DVC_YAML_DATA,
+    )
+
+    context = Context(CONTEXT_DATA)
+    assert resolve(TEMPLATED_DVC_YAML_DATA, context) == RESOLVED_DVC_YAML_DATA
diff --git a/tests/unit/test_stage_resolver.py b/tests/unit/test_stage_resolver.py
new file mode 100644
index 000000000..011199b13
--- /dev/null
+++ b/tests/unit/test_stage_resolver.py
@@ -0,0 +1,39 @@
+from dvc.parsing import DataResolver
+
+TEMPLATED_DVC_YAML_DATA = {
+    "stages": {
+        "stage1": {
+            "cmd": "python script.py ${dict.foo} --out ${dict.bar}",
+            "outs": ["${dict.bar}"],
+            "deps": ["${dict.foo}"],
+            "params": ["${list.0}", "${list.1}"],
+            "frozen": "${freeze}",
+        },
+        "stage2": {"cmd": "echo ${dict.foo} ${dict.bar}"},
+    }
+}
+
+CONTEXT_DATA = {
+    "dict": {"foo": "foo", "bar": "bar"},
+    "list": ["param1", "param2"],
+    "freeze": True,
+}
+
+RESOLVED_DVC_YAML_DATA = {
+    "stages": {
+        "stage1": {
+            "cmd": "python script.py foo --out bar",
+            "outs": ["bar"],
+            "deps": ["foo"],
+            "params": ["param1", "param2"],
+            "frozen": True,
+        },
+        "stage2": {"cmd": "echo foo bar"},
+    }
+}
+
+
+def test_resolver():
+    resolver = DataResolver(TEMPLATED_DVC_YAML_DATA)
+    resolver.context.data = CONTEXT_DATA
+    assert resolver.resolve() == RESOLVED_DVC_YAML_DATA