Fix merge_update not converting data to Node object (#4846)

We wrap every data to one of the `Node` object. But, since we were
passing `self.data` instead of `self`, it was skipping the conversion
and instead was just trying to set data as-is, which would have make
our latter assumptions incorrect, as other codebases assume it to be
a `Node`.

I have added asserts in a few places. Also, I noticed that, on a few
test assumptions, when `resolving`, we were not converting values
back to it's original form (`CtxDict` -> `dict`, `CtxList` -> `list`
and `Value` -> _). So, I have added `node.value` property, as it is
a complement of the above fix.

This caused a ripple effect for the changes in the API as `foreach`
now resolved and changed the data to the original form,
which prevents us from tracking the data later in the "field".
So, `resolve`/`select` got a real `unwrap` option that should
reliably work for all kinds of `Node` objects.

dvc/parsing/__init__.py

@@ -136,7 +136,7 @@ class DataResolver:
             suffix = str(key if key is not DEFAULT_SENTINEL else value)
             return self._resolve_stage(c, f"{name}-{suffix}", in_data)
 
-        iterable = context.resolve(foreach_data)
+        iterable = context.resolve(foreach_data, unwrap=False)
 
         assert isinstance(iterable, (Sequence, Mapping)) and not isinstance(
             iterable, str

dvc/parsing/context.py

@@ -1,4 +1,5 @@
 import os
+from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Mapping, MutableMapping, MutableSequence, Sequence
 from contextlib import contextmanager
@@ -6,7 +7,7 @@ from copy import deepcopy
 from dataclasses import dataclass, field, replace
 from typing import Any, List, Optional, Union
 
-from funcy import identity
+from funcy import identity, rpartial
 
 from dvc.parsing.interpolate import (
     get_expression,
@@ -30,6 +31,7 @@ def _merge(into, update, overwrite):
                     f"Cannot overwrite as key {key} already exists in {into}"
                 )
             into[key] = val
+            assert isinstance(into[key], Node)
 
 
 @dataclass
@@ -55,24 +57,41 @@ def _default_meta():
     return Meta(source=None)
 
 
+class Node:
+    def get_sources(self):
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def value(self):
+        pass
+
+
 @dataclass
-class Value:
-    value: Any
+class Value(Node):
+    _value: Any
     meta: Meta = field(
         compare=False, default_factory=_default_meta, repr=False
     )
 
     def __repr__(self):
-        return repr(self.value)
+        return repr(self._value)
 
     def __str__(self) -> str:
-        return str(self.value)
+        return str(self._value)
 
     def get_sources(self):
         return {self.meta.source: self.meta.path()}
 
+    @property
+    def value(self):
+        return self._value
 
-class Container:
+
+PRIMITIVES = (int, float, str, bytes, bool)
+
+
+class Container(Node, ABC):
     meta: Meta
     data: Union[list, dict]
     _key_transform = staticmethod(identity)
@@ -82,9 +101,9 @@ class Container:
 
     def _convert(self, key, value):
         meta = Meta.update_path(self.meta, key)
-        if value is None or isinstance(value, (int, float, str, bytes, bool)):
+        if value is None or isinstance(value, PRIMITIVES):
             return Value(value, meta=meta)
-        elif isinstance(value, (CtxList, CtxDict, Value)):
+        elif isinstance(value, Node):
             return value
         elif isinstance(value, (list, dict)):
             container = CtxDict if isinstance(value, dict) else CtxList
@@ -150,6 +169,10 @@ class CtxList(Container, MutableSequence):
     def get_sources(self):
         return {self.meta.source: self.meta.path()}
 
+    @property
+    def value(self):
+        return [node.value for node in self]
+
 
 class CtxDict(Container, MutableMapping):
     def __init__(self, mapping: Mapping = None, meta: Meta = None, **kwargs):
@@ -169,7 +192,11 @@ class CtxDict(Container, MutableMapping):
 
     def merge_update(self, *args, overwrite=False):
         for d in args:
-            _merge(self.data, d, overwrite=overwrite)
+            _merge(self, d, overwrite=overwrite)
+
+    @property
+    def value(self):
+        return {key: node.value for key, node in self.items()}
 
 
 class Context(CtxDict):
@@ -205,10 +232,21 @@ class Context(CtxDict):
     def select(
         self, key: str, unwrap=False
     ):  # pylint: disable=arguments-differ
-        # NOTE: `unwrap` default value is different from `resolve_str`
+        """Select the item using key, similar to `__getitem__`
+           but can track the usage of the data on interpolation
+           as well and can get from nested data structure by using
+           "." separated key (eg: "key1.key2.key3")
+
+        Args:
+            key: key to select value from
+            unwrap: Convert CtxList/CtxDict/Value items to it's original data
+                    Defaults to False. Note that the default is different from
+                    `resolve`.
+        """
         node = super().select(key)
         self._track_data(node)
-        return node.value if isinstance(node, Value) and unwrap else node
+        assert isinstance(node, Node)
+        return node.value if unwrap else node
 
     @classmethod
     def load_from(cls, tree, file: str) -> "Context":
@@ -243,20 +281,27 @@ class Context(CtxDict):
             self._check_interpolation_collection(key, value)
         self[key] = value
 
-    def resolve(self, src):
+    def resolve(self, src, unwrap=True):
         """Recursively resolves interpolation and returns resolved data.
 
+        Args:
+            src: Data (str/list/dict etc.) to resolve
+            unwrap: Unwrap CtxDict/CtxList/Value to it's original data if
+                    inside `src`. Defaults to True.
+
         >>> c = Context({"three": 3})
         >>> c.resolve({"lst": [1, 2, "${three}"]})
         {'lst': [1, 2, 3]}
         """
         Seq = (list, tuple, set)
+
+        resolve = rpartial(self.resolve, unwrap)
         if isinstance(src, Mapping):
-            return {key: self.resolve(value) for key, value in src.items()}
+            return {key: resolve(value) for key, value in src.items()}
         elif isinstance(src, Seq):
-            return type(src)(map(self.resolve, src))
+            return type(src)(map(resolve, src))
         elif isinstance(src, str):
-            return self.resolve_str(src)
+            return self.resolve_str(src, unwrap=unwrap)
         return src
 
     def resolve_str(self, src: str, unwrap=True):

tests/func/test_stage_resolver.py

@@ -6,6 +6,7 @@ import pytest
 
 from dvc.dependency import _merge_params
 from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver
+from dvc.parsing.context import Node
 from dvc.path_info import PathInfo
 from dvc.utils.serialize import dump_json, dump_yaml
 
@@ -44,9 +45,19 @@ RESOLVED_DVC_YAML_DATA = {
 }
 
 
+def recurse_not_a_node(d):
+    assert not isinstance(d, Node)
+    if isinstance(d, (list, dict)):
+        iterable = d if isinstance(d, list) else d.values()
+        for item in iterable:
+            assert recurse_not_a_node(item)
+    return True
+
+
 def assert_stage_equal(d1, d2):
     """Keeps the params section in order, and then checks for equality."""
     for d in [d1, d2]:
+        assert recurse_not_a_node(d)
         for _, stage_d in d.get("stages", {}).items():
             params = _merge_params(stage_d.get("params", []))
             for k in params:
@@ -248,12 +259,15 @@ def test_simple_foreach_loop(tmp_dir, dvc):
     }
 
     resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert resolver.resolve() == {
-        "stages": {
-            f"build-{item}": {"cmd": f"python script.py {item}"}
-            for item in iterable
-        }
-    }
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                f"build-{item}": {"cmd": f"python script.py {item}"}
+                for item in iterable
+            }
+        },
+    )
 
 
 def test_foreach_loop_dict(tmp_dir, dvc):
@@ -268,12 +282,15 @@ def test_foreach_loop_dict(tmp_dir, dvc):
     }
 
     resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert resolver.resolve() == {
-        "stages": {
-            f"build-{key}": {"cmd": f"python script.py {item['thresh']}"}
-            for key, item in iterable["models"].items()
-        }
-    }
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                f"build-{key}": {"cmd": f"python script.py {item['thresh']}"}
+                for key, item in iterable["models"].items()
+            }
+        },
+    )
 
 
 def test_foreach_loop_templatized(tmp_dir, dvc):
@@ -319,14 +336,17 @@ def test_set(tmp_dir, dvc, value):
         }
     }
     resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert resolver.resolve() == {
-        "stages": {
-            "build": {
-                "cmd": f"python script.py --thresh {value}",
-                "always_changed": value,
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                "build": {
+                    "cmd": f"python script.py --thresh {value}",
+                    "always_changed": value,
+                }
             }
-        }
-    }
+        },
+    )
 
 
 @pytest.mark.parametrize(
@@ -343,11 +363,14 @@ def test_coll(tmp_dir, dvc, coll):
         }
     }
     resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert resolver.resolve() == {
-        "stages": {
-            "build": {"cmd": "python script.py --thresh 10", "outs": coll}
-        }
-    }
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                "build": {"cmd": "python script.py --thresh 10", "outs": coll}
+            }
+        },
+    )
 
 
 def test_set_with_foreach(tmp_dir, dvc):
@@ -362,12 +385,15 @@ def test_set_with_foreach(tmp_dir, dvc):
         }
     }
     resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert resolver.resolve() == {
-        "stages": {
-            f"build-{item}": {"cmd": f"command --value {item}"}
-            for item in items
-        }
-    }
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                f"build-{item}": {"cmd": f"command --value {item}"}
+                for item in items
+            }
+        },
+    )
 
 
 def test_set_with_foreach_and_on_stage_definition(tmp_dir, dvc):
@@ -388,15 +414,18 @@ def test_set_with_foreach_and_on_stage_definition(tmp_dir, dvc):
         },
     }
     resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert resolver.resolve() == {
-        "stages": {
-            "build-us": {
-                "cmd": "command --value 10",
-                "params": [{"params.json": ["models.us.thresh"]}],
-            },
-            "build-gb": {
-                "cmd": "command --value 15",
-                "params": [{"params.json": ["models.gb.thresh"]}],
-            },
-        }
-    }
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                "build-us": {
+                    "cmd": "command --value 10",
+                    "params": [{"params.json": ["models.us.thresh"]}],
+                },
+                "build-gb": {
+                    "cmd": "command --value 15",
+                    "params": [{"params.json": ["models.gb.thresh"]}],
+                },
+            }
+        },
+    )

tests/unit/test_context.py

@@ -6,6 +6,7 @@ import pytest
 from dvc.parsing.context import Context, CtxDict, CtxList, Value
 from dvc.tree.local import LocalTree
 from dvc.utils.serialize import dump_yaml
+from tests.func.test_stage_resolver import recurse_not_a_node
 
 
 def test_context():
@@ -145,6 +146,22 @@ def test_select():
         assert context.select(f"lst.{i}.bar{i}") == Value(f"bar{i}")
 
 
+def test_select_unwrap():
+    context = Context({"dct": {"foo": "bar"}}, lst=[1, 2, 3], foo="foo")
+
+    assert context.select("dct.foo", unwrap=True) == "bar"
+    assert context.select("lst.0", unwrap=True) == 1
+    assert context.select("foo", unwrap=True) == "foo"
+
+    node = context.select("dct", unwrap=True)
+    assert isinstance(node, dict) and recurse_not_a_node(node)
+    assert node == {"foo": "bar"}
+
+    node = context.select("lst", unwrap=True)
+    assert isinstance(node, list) and recurse_not_a_node(node)
+    assert node == [1, 2, 3]
+
+
 def test_merge_dict():
     d1 = {"Train": {"us": {"lr": 10}}}
     d2 = {"Train": {"us": {"layers": 100}}}
@@ -160,7 +177,11 @@ def test_merge_dict():
         c1.merge_update({"Train": {"us": {"lr": 15}}})
 
     c1.merge_update({"Train": {"us": {"lr": 15}}}, overwrite=True)
-    assert c1.select("Train.us") == CtxDict(lr=15, layers=100)
+    node = c1.select("Train.us")
+    assert node == {"lr": 15, "layers": 100}
+    assert isinstance(node, CtxDict)
+    assert node["lr"] == Value(15)
+    assert node["layers"] == Value(100)
 
 
 def test_merge_list():
@@ -171,7 +192,10 @@ def test_merge_list():
 
     # lists are never merged
     c1.merge_update({"lst": [10, 11, 12]}, overwrite=True)
-    assert c1.select("lst") == [10, 11, 12]
+    node = c1.select("lst")
+    assert node == [10, 11, 12]
+    assert isinstance(node, CtxList)
+    assert node[0] == Value(10)
 
 
 def test_overwrite_with_setitem():
@@ -319,3 +343,30 @@ def test_track_from_multiple_files(tmp_dir):
         assert not key_tracked(path1, "Train.us.layers") and key_tracked(
             path2, "Train.us.layers"
         )
+
+
+def test_node_value():
+    d = {"dct": {"foo": "bar"}, "lst": [1, 2, 3], "foo": "foo"}
+    context = Context(d)
+    assert isinstance(context, (Context, CtxDict))
+    assert isinstance(context["dct"], CtxDict)
+    assert isinstance(context["lst"], CtxList)
+    assert isinstance(context["foo"], Value)
+    assert isinstance(context["dct"]["foo"], Value)
+    assert isinstance(context["lst"][0], Value)
+
+    assert context.value == d
+    assert recurse_not_a_node(context.value)
+    assert isinstance(context.value["dct"], dict)
+    assert isinstance(context.value["lst"], list)
+    assert isinstance(context.value["foo"], str)
+    assert isinstance(context.value["dct"]["foo"], str)
+    assert isinstance(context.value["lst"][0], int)
+
+    assert isinstance(context["dct"].value, dict)
+    assert context["dct"]["foo"].value == "bar"
+
+    assert isinstance(context["lst"].value, list)
+    assert context["lst"][1].value == 2
+
+    assert context["foo"].value == "foo"

tests/unit/test_interpolate.py

@@ -3,6 +3,7 @@ from math import inf, pi
 import pytest
 
 from dvc.parsing import Context
+from tests.func.test_stage_resolver import recurse_not_a_node
 
 
 @pytest.mark.parametrize(
@@ -97,4 +98,6 @@ def test_resolve_collection():
     )
 
     context = Context(CONTEXT_DATA)
-    assert context.resolve(TEMPLATED_DVC_YAML_DATA) == RESOLVED_DVC_YAML_DATA
+    resolved = context.resolve(TEMPLATED_DVC_YAML_DATA)
+    assert resolved == RESOLVED_DVC_YAML_DATA
+    assert recurse_not_a_node(resolved)