refactor resolver and optimize (#5077)

minor refactors, code comments

Rename interim to definitions

provide location of entry during init for definition

small fixups to set_temporarily

Refactor loop to use dict comprehension

Provide same interface for definitions

Make it work for composite lists

Check for sequence rather than a list

refactor tests

Fixes #4955
Fixes #4957
Fixes #4963

dvc/parsing/__init__.py

@@ -1,12 +1,9 @@
 import logging
-import os
-from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from copy import deepcopy
-from itertools import starmap
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, Union
 
-from funcy import join, lfilter
+from funcy import cached_property, collecting, first, isa, join, reraise
 
 from dvc.dependency.param import ParamsDependency
 from dvc.exceptions import DvcException
@@ -19,9 +16,14 @@ from .context import (
     ContextError,
     KeyNotInContext,
     MergeError,
-    Meta,
     Node,
-    SetError,
+    VarsAlreadyLoaded,
+)
+from .interpolate import (
+    check_recursive_parse_errors,
+    is_interpolated_string,
+    recurse,
+    to_str,
 )
 
 if TYPE_CHECKING:
@@ -36,18 +38,17 @@ DEFAULT_PARAMS_FILE = ParamsDependency.DEFAULT_PARAMS_FILE
 PARAMS_KWD = "params"
 FOREACH_KWD = "foreach"
 DO_KWD = "do"
-SET_KWD = "set"
 
 DEFAULT_SENTINEL = object()
 
 JOIN = "@"
 
 
-class VarsAlreadyLoaded(DvcException):
+class ResolveError(DvcException):
     pass
 
 
-class ResolveError(DvcException):
+class EntryNotFound(DvcException):
     pass
 
 
@@ -75,158 +76,144 @@ def _reraise_err(exc_cls, *args, from_exc=None):
     raise err
 
 
+def check_syntax_errors(
+    definition: dict, name: str, path: str, where: str = "stages"
+):
+    for key, d in definition.items():
+        try:
+            check_recursive_parse_errors(d)
+        except ParseError as exc:
+            format_and_raise(exc, f"'{where}.{name}.{key}'", path)
+
+
+def is_map_or_seq(data):
+    _is_map_or_seq = isa(Mapping, Sequence)
+    return not isinstance(data, str) and _is_map_or_seq(data)
+
+
+def split_foreach_name(name):
+    group, *keys = name.rsplit(JOIN, maxsplit=1)
+    return group, first(keys)
+
+
+def check_interpolations(data, where, path):
+    def func(s):
+        if is_interpolated_string(s):
+            raise ResolveError(
+                _format_preamble(f"'{where}'", path)
+                + "interpolating is not allowed"
+            )
+
+    return recurse(func)(data)
+
+
+Definition = Union["ForeachDefinition", "EntryDefinition"]
+DictStr = Dict[str, Any]
+
+
+def make_definition(
+    resolver: "DataResolver", name: str, definition: DictStr, **kwargs
+) -> Definition:
+    args = resolver, resolver.context, name, definition
+    if FOREACH_KWD in definition:
+        return ForeachDefinition(*args, **kwargs)
+    return EntryDefinition(*args, **kwargs)
+
+
 class DataResolver:
     def __init__(self, repo: "Repo", wdir: PathInfo, d: dict):
-        self.data: dict = d
+        self.tree = tree = repo.tree
         self.wdir = wdir
-        self.repo = repo
-        self.tree = self.repo.tree
-        self.imported_files: Dict[str, Optional[List[str]]] = {}
         self.relpath = relpath(self.wdir / "dvc.yaml")
 
-        to_import: PathInfo = wdir / DEFAULT_PARAMS_FILE
-        if self.tree.exists(to_import):
-            self.imported_files = {os.path.abspath(to_import): None}
-            self.global_ctx = Context.load_from(self.tree, to_import)
-        else:
-            self.global_ctx = Context()
-            logger.debug(
-                "%s does not exist, it won't be used in parametrization",
-                to_import,
-            )
-
         vars_ = d.get(VARS_KWD, [])
+        check_interpolations(vars_, VARS_KWD, self.relpath)
+        self.context: Context = Context()
+
         try:
-            self.load_from_vars(
-                self.global_ctx, vars_, wdir, skip_imports=self.imported_files
-            )
-        except (ContextError, VarsAlreadyLoaded) as exc:
+            args = tree, vars_, wdir  # load from `vars` section
+            self.context.load_from_vars(*args, default=DEFAULT_PARAMS_FILE)
+        except ContextError as exc:
             format_and_raise(exc, "'vars'", self.relpath)
 
+        # we wrap the definitions into ForeachDefinition and EntryDefinition,
+        # that helps us to optimize, cache and selectively load each one of
+        # them as we need, and simplify all of this DSL/parsing logic.
+        # we use `tracked_vars` to keep a dictionary of used variables
+        # by the interpolated entries.
         self.tracked_vars: Dict[str, Mapping] = {}
 
-    @staticmethod
-    def check_loaded(path, item, keys, skip_imports):
-        if not keys and isinstance(skip_imports[path], list):
-            raise VarsAlreadyLoaded(
-                f"cannot load '{item}' as it's partially loaded already"
-            )
-        elif keys and skip_imports[path] is None:
-            raise VarsAlreadyLoaded(
-                f"cannot partially load '{item}' as it's already loaded."
-            )
-        elif keys and isinstance(skip_imports[path], list):
-            if not set(keys).isdisjoint(set(skip_imports[path])):
-                raise VarsAlreadyLoaded(
-                    f"cannot load '{item}' as it's partially loaded already"
-                )
+        stages_data = d.get(STAGES_KWD, {})
+        self.definitions: Dict[str, Definition] = {
+            name: make_definition(self, name, definition)
+            for name, definition in stages_data.items()
+        }
 
-    def load_from_vars(
-        self,
-        context: "Context",
-        vars_: List,
-        wdir: PathInfo,
-        skip_imports: Dict[str, Optional[List[str]]],
-        stage_name: str = None,
-    ):
-        stage_name = stage_name or ""
-        for index, item in enumerate(vars_):
-            assert isinstance(item, (str, dict))
-            if isinstance(item, str):
-                path, _, keys_str = item.partition(":")
-                keys = lfilter(bool, keys_str.split(","))
+    def resolve_one(self, name):
+        group, key = split_foreach_name(name)
 
-                path_info = wdir / path
-                path = os.path.abspath(path_info)
+        if not self._has_group_and_key(group, key):
+            raise EntryNotFound(f"Could not find '{name}'")
 
-                if path in skip_imports:
-                    if not keys and skip_imports[path] is None:
-                        # allow specifying complete filepath multiple times
-                        continue
-                    self.check_loaded(path, item, keys, skip_imports)
-
-                context.merge_from(self.tree, path_info, select_keys=keys)
-                skip_imports[path] = keys if keys else None
-            else:
-                joiner = "." if stage_name else ""
-                meta = Meta(source=f"{stage_name}{joiner}vars[{index}]")
-                context.merge_update(Context(item, meta=meta))
-
-    def _resolve_entry(self, name: str, definition):
-        context = Context.clone(self.global_ctx)
-        if FOREACH_KWD in definition:
-            assert DO_KWD in definition
-            self.set_context_from(
-                context, definition.get(SET_KWD, {}), source=[name, "set"]
-            )
-            return self._foreach(
-                context, name, definition[FOREACH_KWD], definition[DO_KWD]
-            )
-
-        try:
-            return self._resolve_stage(context, name, definition)
-        except ContextError as exc:
-            format_and_raise(exc, f"stage '{name}'", self.relpath)
+        # all of the checks for `key` not being None for `ForeachDefinition`
+        # and/or `group` not existing in the `interim`, etc. should be
+        # handled by the `self.has_key()` above.
+        definition = self.definitions[group]
+        if isinstance(definition, EntryDefinition):
+            return definition.resolve()
+        return definition.resolve_one(key)
 
     def resolve(self):
-        stages = self.data.get(STAGES_KWD, {})
-        data = join(starmap(self._resolve_entry, stages.items())) or {}
-        logger.trace("Resolved dvc.yaml:\n%s", data)
+        """Used for testing purposes, otherwise use resolve_one()."""
+        data = join(map(self.resolve_one, self.get_keys()))
+        logger.trace(  # type: ignore[attr-defined]
+            "Resolved dvc.yaml:\n%s", data
+        )
         return {STAGES_KWD: data}
 
-    def _resolve_stage(self, context: Context, name: str, definition) -> dict:
-        definition = deepcopy(definition)
-        self.set_context_from(
-            context, definition.pop(SET_KWD, {}), source=[name, "set"]
-        )
-        wdir = self._resolve_wdir(context, name, definition.get(WDIR_KWD))
-        if self.wdir != wdir:
-            logger.debug(
-                "Stage %s has different wdir than dvc.yaml file", name
-            )
+    def has_key(self, key):
+        return self._has_group_and_key(*split_foreach_name(key))
 
-        vars_ = definition.pop(VARS_KWD, [])
-        # FIXME: Should `vars` be templatized?
+    def _has_group_and_key(self, group, key=None):
         try:
-            self.load_from_vars(
-                context,
-                vars_,
-                wdir,
-                skip_imports=deepcopy(self.imported_files),
-                stage_name=name,
-            )
-        except VarsAlreadyLoaded as exc:
-            format_and_raise(exc, f"'stages.{name}.vars'", self.relpath)
+            definition = self.definitions[group]
+        except KeyError:
+            return False
 
-        logger.trace(  # type: ignore[attr-defined]
-            "Context during resolution of stage %s:\n%s", name, context
-        )
+        if key:
+            return isinstance(
+                definition, ForeachDefinition
+            ) and definition.has_member(key)
+        return not isinstance(definition, ForeachDefinition)
 
-        with context.track():
-            resolved = {}
-            for key, value in definition.items():
-                # NOTE: we do not pop "wdir", and resolve it again
-                # this does not affect anything and is done to try to
-                # track the source of `wdir` interpolation.
-                # This works because of the side-effect that we do not
-                # allow overwriting and/or str interpolating complex objects.
-                # Fix if/when those assumptions are no longer valid.
-                try:
-                    resolved[key] = context.resolve(value)
-                except (ParseError, KeyNotInContext) as exc:
-                    format_and_raise(
-                        exc, f"'stages.{name}.{key}'", self.relpath
-                    )
+    @collecting
+    def get_keys(self):
+        for name, definition in self.definitions.items():
+            if isinstance(definition, ForeachDefinition):
+                yield from definition.get_generated_names()
+                continue
+            yield name
 
-        self.tracked_vars[name] = self._get_vars(context)
-        return {name: resolved}
+    def track_vars(self, name, vars_):
+        self.tracked_vars[name] = vars_
 
-    def _get_vars(self, context: Context):
-        tracked: Dict[str, dict] = defaultdict(dict)
-        for path, vars_ in context.tracked.items():
-            for var in vars_:
-                tracked[path][var] = context.select(var, unwrap=True)
-        return tracked
+
+class EntryDefinition:
+    def __init__(
+        self,
+        resolver: DataResolver,
+        context: Context,
+        name: str,
+        definition: dict,
+        where: str = STAGES_KWD,
+    ):
+        self.resolver = resolver
+        self.wdir = self.resolver.wdir
+        self.relpath = self.resolver.relpath
+        self.context = context
+        self.name = name
+        self.definition = definition
+        self.where = where
 
     def _resolve_wdir(
         self, context: Context, name: str, wdir: str = None
@@ -235,75 +222,147 @@ class DataResolver:
             return self.wdir
 
         try:
-            wdir = str(context.resolve_str(wdir, unwrap=True))
+            wdir = to_str(context.resolve_str(wdir, unwrap=True))
         except (ContextError, ParseError) as exc:
-            format_and_raise(exc, f"'stages.{name}.wdir'", self.relpath)
-        return self.wdir / str(wdir)
+            format_and_raise(exc, f"'{self.where}.{name}.wdir'", self.relpath)
+        return self.wdir / wdir
 
-    def _foreach(self, context: Context, name: str, foreach_data, do_data):
-        iterable = self._resolve_foreach_data(context, name, foreach_data)
-        args = (context, name, do_data, iterable)
-        it = (
-            range(len(iterable))
-            if not isinstance(iterable, Mapping)
-            else iterable
-        )
-        gen = (self._each_iter(*args, i) for i in it)
-        return join(gen)
-
-    def _each_iter(self, context: Context, name: str, do_data, iterable, key):
-        value = iterable[key]
-        c = Context.clone(context)
-        suffix = c["item"] = value
-        if isinstance(iterable, Mapping):
-            suffix = c["key"] = key
-
-        generated = f"{name}{JOIN}{suffix}"
+    def resolve(self, **kwargs):
         try:
-            return self._resolve_stage(c, generated, do_data)
+            return self.resolve_stage(**kwargs)
         except ContextError as exc:
-            # pylint: disable=no-member
-            if isinstance(exc, MergeError) and exc.key in self._inserted_keys(
-                iterable
-            ):
-                raise ResolveError(
-                    f"attempted to redefine '{exc.key}' in stage '{generated}'"
-                    " generated through 'foreach'"
-                )
-            format_and_raise(
-                exc, f"stage '{generated}' (gen. from '{name}')", self.relpath
+            format_and_raise(exc, f"stage '{self.name}'", self.relpath)
+
+    def resolve_stage(self, skip_checks=False):
+        context = self.context
+        name = self.name
+        if not skip_checks:
+            # we can check for syntax errors as we go for interpolated entries,
+            # but for foreach-generated ones, once is enough, which it does
+            # that itself. See `ForeachDefinition.do_definition`.
+            check_syntax_errors(self.definition, name, self.relpath)
+
+        # we need to pop vars from generated/evaluated data
+        definition = deepcopy(self.definition)
+
+        wdir = self._resolve_wdir(context, name, definition.get(WDIR_KWD))
+        if self.wdir != wdir:
+            logger.debug(
+                "Stage %s has different wdir than dvc.yaml file", name
             )
 
-    def _resolve_foreach_data(
-        self, context: "Context", name: str, foreach_data
-    ):
+        vars_ = definition.pop(VARS_KWD, [])
+        # FIXME: Should `vars` be templatized?
+        check_interpolations(vars_, f"{self.where}.{name}.vars", self.relpath)
+        if vars_:
+            # Optimization: Lookahead if it has any vars, if it does not, we
+            # don't need to clone them.
+            context = Context.clone(context)
+
         try:
-            iterable = context.resolve(foreach_data, unwrap=False)
-        except (ContextError, ParseError) as exc:
-            format_and_raise(exc, f"'stages.{name}.foreach'", self.relpath)
-        if isinstance(iterable, str) or not isinstance(
-            iterable, (Sequence, Mapping)
-        ):
-            raise ResolveError(
-                f"failed to resolve 'stages.{name}.foreach'"
-                f" in '{self.relpath}': expected list/dictionary, got "
-                + type(
-                    iterable.value if isinstance(iterable, Node) else iterable
-                ).__name__
+            tree = self.resolver.tree
+            context.load_from_vars(tree, vars_, wdir, stage_name=name)
+        except VarsAlreadyLoaded as exc:
+            format_and_raise(exc, f"'{self.where}.{name}.vars'", self.relpath)
+
+        logger.trace(  # type: ignore[attr-defined]
+            "Context during resolution of stage %s:\n%s", name, context
+        )
+
+        with context.track() as tracked_data:
+            # NOTE: we do not pop "wdir", and resolve it again
+            # this does not affect anything and is done to try to
+            # track the source of `wdir` interpolation.
+            # This works because of the side-effect that we do not
+            # allow overwriting and/or str interpolating complex objects.
+            # Fix if/when those assumptions are no longer valid.
+            resolved = {
+                key: self._resolve(context, value, key, skip_checks)
+                for key, value in definition.items()
+            }
+
+        self.resolver.track_vars(name, tracked_data)
+        return {name: resolved}
+
+    def _resolve(self, context, value, key, skip_checks):
+        try:
+            return context.resolve(
+                value, skip_interpolation_checks=skip_checks
+            )
+        except (ParseError, KeyNotInContext) as exc:
+            format_and_raise(
+                exc, f"'{self.where}.{self.name}.{key}'", self.relpath
             )
 
-        warn_for = [k for k in self._inserted_keys(iterable) if k in context]
+
+class ForeachDefinition:
+    def __init__(
+        self,
+        resolver: DataResolver,
+        context: Context,
+        name: str,
+        definition: dict,
+        where: str = STAGES_KWD,
+    ):
+        self.resolver = resolver
+        self.relpath = self.resolver.relpath
+        self.context = context
+        self.name = name
+
+        assert DO_KWD in definition
+        self.foreach_data = definition[FOREACH_KWD]
+        self._do_definition = definition[DO_KWD]
+
+        self.where = where
+
+    @cached_property
+    def do_definition(self):
+        # optimization: check for syntax errors only once for `foreach` stages
+        check_syntax_errors(self._do_definition, self.name, self.relpath)
+        return self._do_definition
+
+    @cached_property
+    def resolved_iterable(self):
+        return self._resolve_foreach_data()
+
+    def _resolve_foreach_data(self):
+        try:
+            iterable = self.context.resolve(self.foreach_data, unwrap=False)
+        except (ContextError, ParseError) as exc:
+            format_and_raise(
+                exc, f"'{self.where}.{self.name}.foreach'", self.relpath
+            )
+
+        # foreach data can be a resolved dictionary/list.
+        self._check_is_map_or_seq(iterable)
+        # foreach stages will have `item` and `key` added to the context
+        # so, we better warn them if they have them already in the context
+        # from the global vars. We could add them in `set_temporarily`, but
+        # that'd make it display for each iteration.
+        self._warn_if_overwriting(self._inserted_keys(iterable))
+        return iterable
+
+    def _check_is_map_or_seq(self, iterable):
+        if not is_map_or_seq(iterable):
+            node = iterable.value if isinstance(iterable, Node) else iterable
+            typ = type(node).__name__
+            raise ResolveError(
+                f"failed to resolve '{self.where}.{self.name}.foreach'"
+                f" in '{self.relpath}': expected list/dictionary, got " + typ
+            )
+
+    def _warn_if_overwriting(self, keys):
+        warn_for = [k for k in keys if k in self.context]
         if warn_for:
+            linking_verb = "is" if len(warn_for) == 1 else "are"
             logger.warning(
                 "%s %s already specified, "
                 "will be overwritten for stages generated from '%s'",
                 " and ".join(warn_for),
-                "is" if len(warn_for) == 1 else "are",
-                name,
+                linking_verb,
+                self.name,
             )
 
-        return iterable
-
     @staticmethod
     def _inserted_keys(iterable):
         keys = ["item"]
@@ -311,12 +370,74 @@ class DataResolver:
             keys.append("key")
         return keys
 
-    @classmethod
-    def set_context_from(cls, context: Context, to_set, source=None):
-        # NOTE: `set` is disabled in the schema
-        try:
-            for key, value in to_set.items():
-                src_set = [*(source or []), key]
-                context.set(key, value, source=".".join(src_set))
-        except SetError as exc:
-            _reraise_err(ResolveError, str(exc), from_exc=exc)
+    @cached_property
+    def normalized_iterable(self):
+        """Convert sequence to Mapping with keys normalized."""
+        iterable = self.resolved_iterable
+        if isinstance(iterable, Mapping):
+            return iterable
+
+        assert isinstance(iterable, Sequence)
+        if any(map(is_map_or_seq, iterable)):
+            # if the list contains composite data, index are the keys
+            return {to_str(idx): value for idx, value in enumerate(iterable)}
+
+        # for simple lists, eg: ["foo", "bar"],  contents are the key itself
+        return {to_str(value): value for value in iterable}
+
+    def has_member(self, key):
+        return key in self.normalized_iterable
+
+    def get_generated_names(self):
+        return list(map(self._generate_name, self.normalized_iterable))
+
+    def _generate_name(self, key):
+        return f"{self.name}{JOIN}{key}"
+
+    def resolve_all(self):
+        return join(map(self.resolve_one, self.normalized_iterable))
+
+    def resolve_one(self, key):
+        return self._each_iter(key)
+
+    def _each_iter(self, key):
+        name = self.name
+        err_message = f"Could not find '{key}' in foreach group '{self.name}'"
+        with reraise(KeyError, EntryNotFound(err_message)):
+            value = self.normalized_iterable[key]
+
+        # NOTE: we need to use resolved iterable/foreach-data,
+        # not the normalized ones to figure out whether to make item/key
+        # available
+        inserted = self._inserted_keys(self.resolved_iterable)
+        temp_dict = {"item": value}
+        if "key" in inserted:
+            temp_dict["key"] = key
+
+        with self.context.set_temporarily(temp_dict):
+            # optimization: item and key can be removed on __exit__() as they
+            # are top-level values, and are not merged recursively.
+            # This helps us avoid cloning context, which is slower
+            # (increasing the size of the context might increase
+            # the no. of items to be generated which means more cloning,
+            # i.e. quadratic complexity).
+            generated = self._generate_name(key)
+            entry = EntryDefinition(
+                self.resolver, self.context, generated, self.do_definition
+            )
+            try:
+                # optimization: skip checking for syntax errors on each foreach
+                # generated stages. We do it once when accessing do_definition.
+                return entry.resolve_stage(skip_checks=True)
+            except ContextError as exc:
+                # pylint: disable=no-member
+                if isinstance(exc, MergeError) and exc.key in inserted:
+                    raise ResolveError(
+                        f"attempted to redefine '{exc.key}' "
+                        f"in stage '{generated}' generated through 'foreach'"
+                    )
+                format_and_raise(
+                    exc,
+                    f"stage '{generated}' (gen. from '{name}')",
+                    self.relpath,
+                )

dvc/parsing/context.py

@@ -1,3 +1,4 @@
+import logging
 import os
 from abc import ABC, abstractmethod
 from collections import defaultdict
@@ -5,23 +6,24 @@ from collections.abc import Mapping, MutableMapping, MutableSequence, Sequence
 from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field, replace
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, List, Optional, Union
 
-from funcy import identity, rpartial
+from funcy import identity, lfilter
 
 from dvc.exceptions import DvcException
 from dvc.parsing.interpolate import (
-    ParseError,
     get_expression,
     get_matches,
     is_exact_string,
-    is_interpolated_string,
+    normalize_key,
+    recurse,
     str_interpolate,
 )
 from dvc.path_info import PathInfo
 from dvc.utils import relpath
 from dvc.utils.serialize import LOADERS
 
+logger = logging.getLogger(__name__)
 SeqOrMap = Union[Sequence, Mapping]
 
 
@@ -29,10 +31,6 @@ class ContextError(DvcException):
     pass
 
 
-class SetError(ContextError):
-    pass
-
-
 class MergeError(ContextError):
     def __init__(self, key, new, into):
         self.key = key
@@ -61,6 +59,10 @@ class KeyNotInContext(ContextError):
         super().__init__(f"Could not find '{key}'")
 
 
+class VarsAlreadyLoaded(ContextError):
+    pass
+
+
 def _merge(into, update, overwrite):
     for key, val in update.items():
         if isinstance(into.get(key), Mapping) and isinstance(val, Mapping):
@@ -72,6 +74,13 @@ def _merge(into, update, overwrite):
             assert isinstance(into[key], Node)
 
 
+def recurse_not_a_node(data: dict):
+    def func(item):
+        assert not isinstance(item, Node)
+
+    return recurse(func)(data)
+
+
 @dataclass
 class Meta:
     source: Optional[str] = None
@@ -227,6 +236,12 @@ class CtxList(Container, MutableSequence):
     def value(self):
         return [node.value for node in self]
 
+    def __deepcopy__(self, _):
+        # optimization: we don't support overriding a list
+        new = CtxList([])
+        new.data = self.data[:]  # shortcircuting __setitem__
+        return new
+
 
 class CtxDict(Container, MutableMapping):
     def __init__(self, mapping: Mapping = None, meta: Meta = None, **kwargs):
@@ -252,6 +267,14 @@ class CtxDict(Container, MutableMapping):
     def value(self):
         return {key: node.value for key, node in self.items()}
 
+    def __deepcopy__(self, _):
+        new = CtxDict()
+        for k, v in self.items():
+            new.data[k] = (
+                deepcopy(v) if isinstance(v, Container) else v
+            )  # shortcircuting __setitem__
+        return new
+
 
 class Context(CtxDict):
     def __init__(self, *args, **kwargs):
@@ -260,13 +283,16 @@ class Context(CtxDict):
         """
         super().__init__(*args, **kwargs)
         self._track = False
-        self._tracked_data = defaultdict(set)
+        self._tracked_data = defaultdict(dict)
+        self.imports = {}
 
     @contextmanager
     def track(self):
         self._track = True
-        yield
+        yield self._tracked_data
+
         self._track = False
+        self._tracked_data = defaultdict(dict)
 
     def _track_data(self, node):
         if not self._track:
@@ -280,11 +306,7 @@ class Context(CtxDict):
                 continue
             params_file = self._tracked_data[source]
             keys = [keys] if isinstance(keys, str) else keys
-            params_file.update(keys)
-
-    @property
-    def tracked(self) -> Dict[str, set]:
-        return self._tracked_data
+            params_file.update({key: node.value for key in keys})
 
     def select(
         self, key: str, unwrap=False
@@ -300,13 +322,14 @@ class Context(CtxDict):
                     Defaults to False. Note that the default is different from
                     `resolve`.
         """
+        key = normalize_key(key)
         try:
             node = super().select(key)
         except ValueError as exc:
             raise KeyNotInContext(key) from exc
 
-        self._track_data(node)
         assert isinstance(node, Node)
+        self._track_data(node)
         return node.value if unwrap else node
 
     @classmethod
@@ -330,48 +353,99 @@ class Context(CtxDict):
                 ) from exc
 
         meta = Meta(source=file, local=False)
-        return cls(data, meta=meta)
+        ctx = cls(data, meta=meta)
+        ctx.imports[os.path.abspath(path)] = select_keys or None
+        return ctx
 
     def merge_from(
-        self, tree, path: PathInfo, overwrite=False, select_keys=None,
+        self, tree, item: str, wdir: PathInfo, overwrite=False,
     ):
-        self.merge_update(
-            Context.load_from(tree, path, select_keys), overwrite=overwrite
-        )
+        path, _, keys_str = item.partition(":")
+        select_keys = lfilter(bool, keys_str.split(","))
+        path_info = wdir / path
+
+        abspath = os.path.abspath(path_info)
+        if abspath in self.imports:
+            if not select_keys and self.imports[abspath] is None:
+                return  # allow specifying complete filepath multiple times
+            self.check_loaded(abspath, item, select_keys)
+
+        ctx = Context.load_from(tree, path_info, select_keys)
+        self.merge_update(ctx, overwrite=overwrite)
+
+        cp = ctx.imports[abspath]
+        if abspath not in self.imports:
+            self.imports[abspath] = cp
+        elif cp:
+            self.imports[abspath].extend(cp)
+
+    def check_loaded(self, path, item, keys):
+        if not keys and isinstance(self.imports[path], list):
+            raise VarsAlreadyLoaded(
+                f"cannot load '{item}' as it's partially loaded already"
+            )
+        elif keys and self.imports[path] is None:
+            raise VarsAlreadyLoaded(
+                f"cannot partially load '{item}' as it's already loaded."
+            )
+        elif keys and isinstance(self.imports[path], list):
+            if not set(keys).isdisjoint(set(self.imports[path])):
+                raise VarsAlreadyLoaded(
+                    f"cannot load '{item}' as it's partially loaded already"
+                )
+
+    def load_from_vars(
+        self,
+        tree,
+        vars_: List,
+        wdir: PathInfo,
+        stage_name: str = None,
+        default: str = None,
+    ):
+        if default:
+            to_import = wdir / default
+            if tree.exists(to_import):
+                self.merge_from(tree, default, wdir)
+            else:
+                msg = "%s does not exist, it won't be used in parametrization"
+                logger.trace(msg, to_import)  # type: ignore[attr-defined]
+
+        stage_name = stage_name or ""
+        for index, item in enumerate(vars_):
+            assert isinstance(item, (str, dict))
+            if isinstance(item, str):
+                self.merge_from(tree, item, wdir)
+            else:
+                joiner = "." if stage_name else ""
+                meta = Meta(source=f"{stage_name}{joiner}vars[{index}]")
+                self.merge_update(Context(item, meta=meta))
+
+    def __deepcopy__(self, _):
+        new = Context(super().__deepcopy__(_))
+        new.meta = deepcopy(self.meta)
+        new.imports = deepcopy(self.imports)
+        return new
 
     @classmethod
     def clone(cls, ctx: "Context") -> "Context":
         """Clones given context."""
-        return cls(deepcopy(ctx.data))
+        return deepcopy(ctx)
 
-    def set(self, key, value, source=None):
-        """
-        Sets a value, either non-interpolated values to a key,
-        or an interpolated string after resolving it.
+    @contextmanager
+    def set_temporarily(self, to_set):
+        non_existing = frozenset(to_set.keys() - self.keys())
+        prev = {key: self[key] for key in to_set if key not in non_existing}
+        to_set = CtxDict(to_set)
+        self.update(to_set)
 
-        >>> c = Context({"foo": "foo", "bar": [1, 2], "lorem": {"a": "z"}})
-        >>> c.set("foobar", "${bar}")
-        >>> c
-        {'foo': 'foo', 'bar': [1, 2], 'lorem': {'a': 'z'}, 'foobar': [1, 2]}
-        """
         try:
-            self._set(key, value, source)
-        except (ParseError, ValueError, ContextError) as exc:
-            sp = "\n" if isinstance(exc, ParseError) else " "
-            raise SetError(f"Failed to set '{source}':{sp}{str(exc)}") from exc
+            yield
+        finally:
+            self.update(prev)
+            for key in non_existing:
+                self.data.pop(key, None)
 
-    def _set(self, key, value, source):
-        if key in self:
-            raise ValueError(f"Cannot set '{key}', key already exists")
-        if isinstance(value, str):
-            self._check_joined_with_interpolation(key, value)
-            value = self.resolve_str(value, unwrap=False)
-        elif isinstance(value, (Sequence, Mapping)):
-            self._check_not_nested_collection(key, value)
-            self._check_interpolation_collection(key, value)
-        self[key] = self._convert_with_meta(value, Meta(source=source))
-
-    def resolve(self, src, unwrap=True):
+    def resolve(self, src, unwrap=True, skip_interpolation_checks=False):
         """Recursively resolves interpolation and returns resolved data.
 
         Args:
@@ -383,18 +457,12 @@ class Context(CtxDict):
         >>> c.resolve({"lst": [1, 2, "${three}"]})
         {'lst': [1, 2, 3]}
         """
-        Seq = (list, tuple, set)
+        func = recurse(self.resolve_str)
+        return func(src, unwrap, skip_interpolation_checks)
 
-        resolve = rpartial(self.resolve, unwrap)
-        if isinstance(src, Mapping):
-            return dict(map(resolve, kv) for kv in src.items())
-        elif isinstance(src, Seq):
-            return type(src)(map(resolve, src))
-        elif isinstance(src, str):
-            return self.resolve_str(src, unwrap=unwrap)
-        return src
-
-    def resolve_str(self, src: str, unwrap=True):
+    def resolve_str(
+        self, src: str, unwrap=True, skip_interpolation_checks=False
+    ):
         """Resolves interpolated string to it's original value,
         or in case of multiple interpolations, a combined string.
 
@@ -402,47 +470,20 @@ class Context(CtxDict):
         >>> c.resolve_str("${enabled}")
         True
         >>> c.resolve_str("enabled? ${enabled}")
-        'enabled? True'
+        'enabled? true'
         """
         matches = get_matches(src)
         if is_exact_string(src, matches):
             # replace "${enabled}", if `enabled` is a boolean, with it's actual
             # value rather than it's string counterparts.
-            expr = get_expression(matches[0])
+            expr = get_expression(
+                matches[0], skip_checks=skip_interpolation_checks
+            )
             return self.select(expr, unwrap=unwrap)
         # but not "${num} days"
-        return str_interpolate(src, matches, self)
-
-    @staticmethod
-    def _check_not_nested_collection(key: str, value: SeqOrMap):
-        values = value.values() if isinstance(value, Mapping) else value
-        has_nested = any(
-            not isinstance(item, str) and isinstance(item, (Mapping, Sequence))
-            for item in values
+        return str_interpolate(
+            src, matches, self, skip_checks=skip_interpolation_checks
         )
-        if has_nested:
-            raise ValueError(f"Cannot set '{key}', has nested dict/list")
-
-    @staticmethod
-    def _check_interpolation_collection(key: str, value: SeqOrMap):
-        values = value.values() if isinstance(value, Mapping) else value
-        interpolated = any(is_interpolated_string(item) for item in values)
-        if interpolated:
-            raise ValueError(
-                f"Cannot set '{key}', "
-                "having interpolation inside "
-                f"'{type(value).__name__}' is not supported."
-            )
-
-    @staticmethod
-    def _check_joined_with_interpolation(key: str, value: str):
-        matches = get_matches(value)
-        if matches and not is_exact_string(value, matches):
-            raise ValueError(
-                f"Cannot set '{key}', "
-                "joining string with interpolated string "
-                "is not supported"
-            )
 
 
 if __name__ == "__main__":

dvc/parsing/interpolate.py

@@ -1,7 +1,9 @@
 import re
 import typing
+from collections.abc import Mapping
 from functools import singledispatch
 
+from funcy import rpartial
 from pyparsing import (
     CharsNotIn,
     ParseException,
@@ -18,6 +20,8 @@ if typing.TYPE_CHECKING:
 
     from .context import Context
 
+BRACE_OPEN = "${"
+BRACE_CLOSE = "}"
 LBRACK = "["
 RBRACK = "]"
 PERIOD = "."
@@ -51,15 +55,23 @@ def get_matches(template: str):
 
 
 def is_interpolated_string(val):
-    return bool(get_matches(val)) if isinstance(val, str) else False
+    return isinstance(val, str) and bool(get_matches(val))
+
+
+def normalize_key(key: str):
+    return key.replace(LBRACK, PERIOD).replace(RBRACK, "")
 
 
 def format_and_raise_parse_error(exc):
     raise ParseError(_format_exc_msg(exc))
 
 
+def embrace(s: str):
+    return BRACE_OPEN + s + BRACE_CLOSE
+
+
 @singledispatch
-def to_str(obj):
+def to_str(obj) -> str:
     return str(obj)
 
 
@@ -72,20 +84,47 @@ def _format_exc_msg(exc: ParseException):
     exc.loc += 2  # 2 because we append `${` at the start of expr below
 
     expr = exc.pstr
-    exc.pstr = "${" + exc.pstr + "}"
+    exc.pstr = embrace(exc.pstr)
     error = ParseException.explain(exc, depth=0)
 
     _, pointer, *explains = error.splitlines()
     pstr = "{brace_open}{expr}{brace_close}".format(
-        brace_open=colorize("${", color="blue"),
+        brace_open=colorize(BRACE_OPEN, color="blue"),
         expr=colorize(expr, color="magenta"),
-        brace_close=colorize("}", color="blue"),
+        brace_close=colorize(BRACE_CLOSE, color="blue"),
     )
     msg = "\n".join(explains)
     pointer = colorize(pointer, color="red")
     return "\n".join([pstr, pointer, colorize(msg, color="red", style="bold")])
 
 
+def recurse(f):
+    Seq = (list, tuple, set)
+
+    def wrapper(data, *args):
+        g = rpartial(wrapper, *args)
+        if isinstance(data, Mapping):
+            return {g(k): g(v) for k, v in data.items()}
+        elif isinstance(data, Seq):
+            return type(data)(map(g, data))
+        elif isinstance(data, str):
+            return f(data, *args)
+        return data
+
+    return wrapper
+
+
+def check_recursive_parse_errors(data: dict):
+    func = recurse(check_expression)
+    return func(data)
+
+
+def check_expression(s: str):
+    matches = get_matches(s)
+    for match in matches:
+        get_expression(match)
+
+
 def parse_expr(s: str):
     try:
         result = parser.parseString(s, parseAll=True)
@@ -97,18 +136,23 @@ def parse_expr(s: str):
     return joined[0]
 
 
-def get_expression(match: "Match"):
+def get_expression(match: "Match", skip_checks: bool = False):
     _, _, inner = match.groups()
-    return parse_expr(inner)
+    return inner if skip_checks else parse_expr(inner)
 
 
-def str_interpolate(template: str, matches: "List[Match]", context: "Context"):
+def str_interpolate(
+    template: str,
+    matches: "List[Match]",
+    context: "Context",
+    skip_checks: bool = False,
+):
     from .context import PRIMITIVES
 
     index, buf = 0, ""
     for match in matches:
         start, end = match.span(0)
-        expr = get_expression(match)
+        expr = get_expression(match, skip_checks=skip_checks)
         value = context.select(expr, unwrap=True)
         if value is not None and not isinstance(value, PRIMITIVES):
             raise ParseError(
@@ -119,7 +163,7 @@ def str_interpolate(template: str, matches: "List[Match]", context: "Context"):
     buf += template[index:]
     # regex already backtracks and avoids any `${` starting with
     # backslashes(`\`). We just need to replace those by `${`.
-    return buf.replace(r"\${", "${")
+    return buf.replace(r"\${", BRACE_OPEN)
 
 
 def is_exact_string(src: str, matches: "List[Match]"):

dvc/stage/loader.py

@@ -3,11 +3,11 @@ from collections.abc import Mapping
 from copy import deepcopy
 from itertools import chain
 
-from funcy import cached_property, get_in, lcat, log_durations, project
+from funcy import cached_property, get_in, lcat, project
 
 from dvc import dependency, output
 from dvc.hash_info import HashInfo
-from dvc.parsing import FOREACH_KWD, JOIN, DataResolver
+from dvc.parsing import FOREACH_KWD, JOIN, DataResolver, EntryNotFound
 from dvc.path_info import PathInfo
 
 from . import PipelineStage, Stage, loads_from
@@ -31,13 +31,6 @@ class StageLoader(Mapping):
         wdir = PathInfo(self.dvcfile.path).parent
         return DataResolver(self.repo, wdir, self.data)
 
-    @cached_property
-    def resolved_data(self):
-        data = self.data
-        with log_durations(logger.debug, "resolving values"):
-            data = self.resolver.resolve()
-        return data.get("stages", {})
-
     @staticmethod
     def fill_from_lock(stage, lock_data=None):
         """Fill values for params, checksums for outs and deps from lock."""
@@ -99,7 +92,9 @@ class StageLoader(Mapping):
         if not name:
             raise StageNameUnspecified(self.dvcfile)
 
-        if name not in self:
+        try:
+            resolved_data = self.resolver.resolve_one(name)
+        except EntryNotFound:
             raise StageNotFound(self.dvcfile, name)
 
         if not self.lockfile_data.get(name):
@@ -107,10 +102,11 @@ class StageLoader(Mapping):
                 "No lock entry found for '%s:%s'", self.dvcfile.relpath, name,
             )
 
+        resolved_stage = resolved_data[name]
         stage = self.load_stage(
             self.dvcfile,
             name,
-            self.resolved_data[name],
+            resolved_stage,
             self.lockfile_data.get(name, {}),
         )
 
@@ -120,19 +116,18 @@ class StageLoader(Mapping):
             stage.raw_data.generated_from = group
 
         stage.raw_data.parametrized = (
-            self.stages_data.get(name, {}) != self.resolved_data[name]
+            self.stages_data.get(name, {}) != resolved_stage
         )
-
         return stage
 
     def __iter__(self):
-        return iter(self.resolved_data)
+        return iter(self.resolver.get_keys())
 
     def __len__(self):
-        return len(self.resolved_data)
+        return len(self.resolver.get_keys())
 
     def __contains__(self, name):
-        return name in self.resolved_data
+        return self.resolver.has_key(name)  # noqa: W601
 
     def is_foreach_generated(self, name: str):
         return (

tests/func/parsing/__init__.py

@@ -0,0 +1,54 @@
+from dvc.parsing import DataResolver, EntryDefinition, ForeachDefinition
+from dvc.parsing.context import Context
+
+TEMPLATED_DVC_YAML_DATA = {
+    "stages": {
+        "stage1": {
+            "cmd": "python script.py ${dict.foo} --out ${dict.bar}",
+            "outs": ["${dict.bar}"],
+            "deps": ["${dict.foo}"],
+            "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"],
+            "frozen": True,
+        },
+        "stage2": {"cmd": "echo foo bar"},
+    }
+}
+
+USED_VARS = {
+    "stage1": {"dict.foo": "foo", "dict.bar": "bar", "freeze": True},
+    "stage2": {"dict.foo": "foo", "dict.bar": "bar"},
+}
+
+
+def make_entry_definition(wdir, name, data, context=None) -> EntryDefinition:
+    return EntryDefinition(
+        DataResolver(wdir.dvc, wdir, {}), context or Context(), name, data
+    )
+
+
+def make_foreach_def(
+    wdir, name, foreach_data, do_data=None, context=None
+) -> ForeachDefinition:
+    return ForeachDefinition(
+        DataResolver(wdir.dvc, wdir, {}),
+        context or Context(),
+        name,
+        {"foreach": foreach_data, "do": do_data or {}},
+    )

tests/func/parsing/test_errors.py

@@ -0,0 +1,313 @@
+"""Negative tests for the parametrization."""
+
+
+import logging
+import re
+
+import pytest
+from funcy import first
+
+from dvc.parsing import ResolveError
+from dvc.parsing.context import Context
+from dvc.parsing.interpolate import embrace
+from dvc.utils.serialize import dump_yaml
+
+from . import make_entry_definition, make_foreach_def
+
+
+def escape_ansi(line):
+    ansi_escape = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]")
+    return ansi_escape.sub("", line)
+
+
+# Tests for the interpolated entries
+
+
+@pytest.mark.parametrize("vars_", ["${file}_params.yaml", {"foo": "${foo}"}])
+def test_vars_interpolation_errors(tmp_dir, dvc, vars_):
+    definition = make_entry_definition(tmp_dir, "build", {"vars": [vars_]})
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve()
+
+    assert (
+        str(exc_info.value)
+        == "failed to parse 'stages.build.vars' in 'dvc.yaml': "
+        "interpolating is not allowed"
+    )
+
+
+def test_failed_to_interpolate(tmp_dir, dvc):
+    context = Context(models={"foo": "bar"})
+    definition = make_entry_definition(
+        tmp_dir, "build", {"cmd": "echo ${models.foo.}"}, context
+    )
+
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve()
+
+    assert escape_ansi(str(exc_info.value)) == (
+        "failed to parse 'stages.build.cmd' in 'dvc.yaml':\n"
+        "${models.foo.}\n"
+        "            ^\n"
+        "ParseException: Expected end of text, found '.'"
+        "  (at char 12), (line:1, col:13)"
+    )
+    assert definition.context == {"models": {"foo": "bar"}}
+
+
+def test_local_vars_params_file_not_exist(tmp_dir, dvc):
+    definition = make_entry_definition(
+        tmp_dir,
+        "build",
+        {"vars": ["not_existing_params.yaml"], "cmd": "echo ${models.foo}"},
+    )
+
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve()
+
+    assert str(exc_info.value) == (
+        "failed to parse stage 'build' in 'dvc.yaml': "
+        "'not_existing_params.yaml' does not exist"
+    )
+    assert not definition.context
+
+
+def test_specified_key_does_not_exist(tmp_dir, dvc):
+    definition = make_entry_definition(
+        tmp_dir,
+        "build",
+        {"cmd": "echo ${models.foobar}"},
+        Context(models={"foo": "foo"}),
+    )
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve()
+
+    assert str(exc_info.value) == (
+        "failed to parse 'stages.build.cmd' in 'dvc.yaml': "
+        "Could not find 'models.foobar'"
+    )
+    assert definition.context == {"models": {"foo": "foo"}}
+
+
+@pytest.mark.parametrize(
+    "wdir, expected_msg",
+    [
+        ("${models[foobar]}", " Could not find 'models.foobar'"),
+        (
+            "${models.foo]}",
+            "\n${models.foo]}\n"
+            "            ^\n"
+            "ParseException: Expected end of text, found ']'"
+            "  (at char 12), (line:1, col:13)",
+        ),
+    ],
+)
+def test_wdir_failed_to_interpolate(tmp_dir, dvc, wdir, expected_msg):
+    definition = make_entry_definition(
+        tmp_dir,
+        "build",
+        {"wdir": wdir, "cmd": "echo ${models.bar}"},
+        Context(models={"bar": "bar"}),
+    )
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve()
+
+    assert escape_ansi(str(exc_info.value)) == (
+        "failed to parse 'stages.build.wdir' in 'dvc.yaml':" + expected_msg
+    )
+    assert definition.context == {"models": {"bar": "bar"}}
+
+
+def test_interpolate_non_string(tmp_dir, dvc):
+    definition = make_entry_definition(
+        tmp_dir, "build", {"cmd": "echo ${models}"}, Context(models={})
+    )
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve()
+
+    assert str(exc_info.value) == (
+        "failed to parse 'stages.build.cmd' in 'dvc.yaml':\n"
+        "Cannot interpolate data of type 'dict'"
+    )
+    assert definition.context == {"models": {}}
+
+
+def test_partial_vars_doesnot_exist(tmp_dir, dvc):
+    dump_yaml("test_params.yaml", {"sub1": "sub1", "sub2": "sub2"})
+
+    definition = make_entry_definition(
+        tmp_dir,
+        "build",
+        {"vars": ["test_params.yaml:sub3"], "cmd": "echo ${sub1} ${sub2}"},
+    )
+
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve()
+
+    assert str(exc_info.value) == (
+        "failed to parse stage 'build' in 'dvc.yaml': "
+        "could not find 'sub3' in 'test_params.yaml'"
+    )
+    assert not definition.context
+
+
+# Tests foreach generated stages and their error messages
+
+
+def test_foreach_data_syntax_error(tmp_dir, dvc):
+    definition = make_foreach_def(tmp_dir, "build", "${syntax.[error}", {})
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve_all()
+
+    assert escape_ansi(str(exc_info.value)) == (
+        "failed to parse 'stages.build.foreach' in 'dvc.yaml':\n"
+        "${syntax.[error}\n"
+        "        ^\n"
+        "ParseException: Expected end of text, found '.'"
+        "  (at char 8), (line:1, col:9)"
+    )
+
+
+@pytest.mark.parametrize("key", ["modelss", "modelss.123"])
+def test_foreach_data_key_does_not_exists(tmp_dir, dvc, key):
+    definition = make_foreach_def(tmp_dir, "build", embrace(key), {})
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve_all()
+    assert str(exc_info.value) == (
+        "failed to parse 'stages.build.foreach' in 'dvc.yaml': "
+        f"Could not find '{key}'"
+    )
+
+
+@pytest.mark.parametrize(
+    "foreach_data", ["${foo}", "${dct.model1}", "${lst.0}", "foobar"],
+)
+def test_foreach_data_expects_list_or_dict(tmp_dir, dvc, foreach_data):
+    context = Context(
+        {"foo": "bar", "dct": {"model1": "a-out"}, "lst": ["foo", "bar"]}
+    )
+    definition = make_foreach_def(tmp_dir, "build", foreach_data, {}, context)
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve_all()
+    assert str(exc_info.value) == (
+        "failed to resolve 'stages.build.foreach' in 'dvc.yaml': "
+        "expected list/dictionary, got str"
+    )
+
+
+@pytest.mark.parametrize(
+    "global_data, where",
+    [
+        ({"item": 10, "key": 10}, "item and key are"),
+        ({"item": 10}, "item is"),
+        ({"key": 5}, "key is"),
+    ],
+)
+def test_foreach_overwriting_item_in_list(
+    tmp_dir, dvc, caplog, global_data, where
+):
+    context = Context(global_data)
+    definition = make_foreach_def(
+        tmp_dir, "build", {"model1": 10, "model2": 5}, {}, context,
+    )
+    with caplog.at_level(logging.WARNING, logger="dvc.parsing"):
+        definition.resolve_all()
+
+    assert caplog.messages == [
+        f"{where} already specified, "
+        "will be overwritten for stages generated from 'build'"
+    ]
+
+
+def test_foreach_do_syntax_errors(tmp_dir, dvc):
+    definition = make_foreach_def(
+        tmp_dir, "build", ["foo", "bar"], {"cmd": "echo ${syntax.[error}"},
+    )
+
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve_all()
+
+    assert escape_ansi(str(exc_info.value)) == (
+        "failed to parse 'stages.build.cmd' in 'dvc.yaml':\n"
+        "${syntax.[error}\n"
+        "        ^\n"
+        "ParseException: Expected end of text, found '.'"
+        "  (at char 8), (line:1, col:9)"
+    )
+
+
+@pytest.mark.parametrize(
+    "key, loc",
+    [
+        (
+            "item.thresh",  # the `thresh` in not available on model2`
+            "[email protected]",
+        ),
+        ("foo.bar", "[email protected]"),  # not available on any stages
+    ],
+)
+def test_foreach_do_definition_item_does_not_exist(tmp_dir, dvc, key, loc):
+    context = Context(foo="bar")
+    definition = make_foreach_def(
+        tmp_dir,
+        "build",
+        [{"thresh": "10"}, {}],
+        {"cmd": embrace(key)},
+        context,
+    )
+
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve_all()
+
+    assert str(exc_info.value) == (
+        f"failed to parse '{loc}' in 'dvc.yaml': Could not find '{key}'"
+    )
+
+    # should have no `item` and `key` even though it failed to resolve.
+    assert context == {"foo": "bar"}
+
+
+@pytest.mark.parametrize(
+    "redefine", [{"item": 5}, {"key": 5}, {"item": 5, "key": 10}],
+)
+@pytest.mark.parametrize("from_file", [True, False])
+def test_item_key_in_generated_stage_vars(tmp_dir, dvc, redefine, from_file):
+    context = Context(foo="bar")
+    vars_ = [redefine]
+    if from_file:
+        dump_yaml("test_params.yaml", redefine)
+        vars_ = ["test_params.yaml"]
+
+    definition = make_foreach_def(
+        tmp_dir,
+        "build",
+        {"model1": {"thresh": "10"}, "model2": {"thresh": 5}},
+        {"vars": vars_, "cmd": "${item}"},
+        context,
+    )
+
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve_all()
+    assert str(exc_info.value) == (
+        f"attempted to redefine '{first(redefine)}' in stage 'build@model1'"
+        " generated through 'foreach'"
+    )
+    assert context == {"foo": "bar"}
+
+
+def test_foreach_wdir_key_does_not_exist(tmp_dir, dvc):
+    definition = make_foreach_def(
+        tmp_dir,
+        "build",
+        "${models}",
+        {"wdir": "${ite}", "cmd": "echo ${item}"},
+        Context(models=["foo", "bar"]),
+    )
+    with pytest.raises(ResolveError) as exc_info:
+        definition.resolve_all()
+    assert (
+        str(exc_info.value)
+        == "failed to parse '[email protected]' in 'dvc.yaml': "
+        "Could not find 'ite'"
+    )
+    assert definition.context == {"models": ["foo", "bar"]}

tests/func/parsing/test_foreach.py

@@ -0,0 +1,419 @@
+"""Testing happy paths for the foreach."""
+import os
+
+import pytest
+
+from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver, ForeachDefinition
+from dvc.parsing.context import Context
+from dvc.utils.serialize import dump_yaml
+
+
+def test_with_simple_list_data(tmp_dir, dvc):
+    """Testing a simple non-nested list as a foreach data"""
+    resolver = DataResolver(dvc, tmp_dir, {})
+
+    context = Context()
+    data = {"foreach": ["foo", "bar", "baz"], "do": {"cmd": "echo ${item}"}}
+    definition = ForeachDefinition(resolver, context, "build", data)
+
+    assert definition.resolve_one("foo") == {"build@foo": {"cmd": "echo foo"}}
+    assert definition.resolve_one("bar") == {"build@bar": {"cmd": "echo bar"}}
+    # check that `foreach` item-key replacement didnot leave any leftovers.
+    assert not context
+    assert not resolver.tracked_vars["build@foo"]
+    assert not resolver.tracked_vars["build@bar"]
+
+
+def test_with_dict_data(tmp_dir, dvc):
+    resolver = DataResolver(dvc, tmp_dir, {})
+    context = Context()
+
+    foreach_data = {"model1": "foo", "model2": "bar"}
+    data = {"foreach": foreach_data, "do": {"cmd": "echo ${key} ${item}"}}
+    definition = ForeachDefinition(resolver, context, "build", data)
+
+    assert definition.resolve_one("model1") == {
+        "build@model1": {"cmd": "echo model1 foo"}
+    }
+    assert definition.resolve_one("model2") == {
+        "build@model2": {"cmd": "echo model2 bar"}
+    }
+
+    # check that `foreach` item-key replacement didnot leave any leftovers.
+    assert not context
+    assert not resolver.tracked_vars["build@model1"]
+    assert not resolver.tracked_vars["build@model2"]
+
+
+def test_with_composite_list(tmp_dir, dvc):
+    resolver = DataResolver(dvc, tmp_dir, {})
+
+    context = Context()
+    foreach_data = [{"thresh": "foo"}, {"thresh": "bar"}]
+    data = {"foreach": foreach_data, "do": {"cmd": "echo ${item.thresh}"}}
+    definition = ForeachDefinition(resolver, context, "build", data)
+
+    assert definition.resolve_one("0") == {"build@0": {"cmd": "echo foo"}}
+    # check that `foreach` item-key replacement didnot leave any leftovers.
+    assert not context
+
+    assert definition.resolve_one("1") == {"build@1": {"cmd": "echo bar"}}
+    assert not context
+    assert not resolver.tracked_vars["build@0"]
+
+
+def test_foreach_interpolated_simple_list(tmp_dir, dvc):
+    foreach_data = ["foo", "bar", "baz"]
+    vars_ = {"models": foreach_data}
+    resolver = DataResolver(dvc, tmp_dir, {"vars": [vars_]})
+    data = {"foreach": "${models}", "do": {"cmd": "echo ${item}"}}
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        "build@foo": {"cmd": "echo foo"},
+        "build@bar": {"cmd": "echo bar"},
+        "build@baz": {"cmd": "echo baz"},
+    }
+    assert resolver.context == {"models": foreach_data}
+    assert not any(item for item in resolver.tracked_vars.values())
+
+
+@pytest.mark.parametrize("foreach_def", ["${item.thresh}", "${item[thresh]}"])
+@pytest.mark.parametrize(
+    "foreach_data, result",
+    [
+        (
+            {"model1": {"thresh": "foo"}, "model2": {"thresh": "bar"}},
+            {
+                "build@model1": {"cmd": "echo foo"},
+                "build@model2": {"cmd": "echo bar"},
+            },
+        ),
+        (
+            [{"thresh": "foo"}, {"thresh": "bar"}],
+            {"build@0": {"cmd": "echo foo"}, "build@1": {"cmd": "echo bar"}},
+        ),
+    ],
+)
+def test_foreach_interpolate_with_composite_data(
+    tmp_dir, dvc, foreach_def, foreach_data, result
+):
+    vars_ = [{"models": foreach_data}]
+    resolver = DataResolver(dvc, tmp_dir, {"vars": vars_})
+    data = {"foreach": "${models}", "do": {"cmd": f"echo {foreach_def}"}}
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == result
+    assert resolver.context == {"models": foreach_data}
+    assert not any(item for item in resolver.tracked_vars.values())
+
+
+def test_params_file_with_dict_tracked(tmp_dir, dvc):
+    foreach_data = {"model1": {"thresh": "foo"}, "model2": {"thresh": "bar"}}
+    params = {"models": foreach_data}
+    dump_yaml("params.yaml", params)
+
+    resolver = DataResolver(dvc, tmp_dir, {})
+    data = {"foreach": "${models}", "do": {"cmd": "echo ${item.thresh}"}}
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        "build@model1": {"cmd": "echo foo"},
+        "build@model2": {"cmd": "echo bar"},
+    }
+    # check that `foreach` item-key replacement didnot leave any leftovers.
+    assert resolver.context == {"models": foreach_data}
+    assert resolver.tracked_vars == {
+        "build@model1": {"params.yaml": {"models.model1.thresh": "foo"}},
+        "build@model2": {"params.yaml": {"models.model2.thresh": "bar"}},
+    }
+
+
+def test_params_file_tracked_for_composite_list(tmp_dir, dvc):
+    foreach_data = [
+        {"thresh": "foo"},
+        {"thresh": "bar"},
+    ]
+    params = {"models": foreach_data}
+    dump_yaml("params.yaml", params)
+
+    resolver = DataResolver(dvc, tmp_dir, {})
+    data = {"foreach": "${models}", "do": {"cmd": "echo ${item.thresh}"}}
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        "build@0": {"cmd": "echo foo"},
+        "build@1": {"cmd": "echo bar"},
+    }
+    assert resolver.context == {"models": foreach_data}
+    assert resolver.tracked_vars == {
+        "build@0": {"params.yaml": {"models.0.thresh": "foo"}},
+        "build@1": {"params.yaml": {"models.1.thresh": "bar"}},
+    }
+
+
+def test_foreach_data_from_nested_vars(tmp_dir, dvc):
+    vars_ = {"models": {"lst": [{"thresh": 10}, {"thresh": 15}]}}
+    resolver = DataResolver(dvc, tmp_dir, {"vars": [vars_]})
+    data = {"foreach": "${models.lst}", "do": {"cmd": "echo ${item.thresh}"}}
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        "build@0": {"cmd": "echo 10"},
+        "build@1": {"cmd": "echo 15"},
+    }
+    assert resolver.context == vars_
+    assert not any(item for item in resolver.tracked_vars.values())
+
+
+def test_foreach_partial_interpolations(tmp_dir, dvc):
+    resolver = DataResolver(dvc, tmp_dir, {"vars": [{"bar": "bar"}]})
+    foreach_data = {"model1": "foo", "model2": "${bar}"}
+    data = {"foreach": foreach_data, "do": {"cmd": "echo ${item}"}}
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        "build@model1": {"cmd": "echo foo"},
+        "build@model2": {"cmd": "echo bar"},
+    }
+    assert resolver.context == {"bar": "bar"}
+    assert not any(item for item in resolver.tracked_vars.values())
+
+
+def test_mixed_vars_for_foreach_data(tmp_dir, dvc):
+    dump_yaml("params.yaml", {"models": {"model1": "foo"}})
+    dump_yaml("test_params.yaml", {"models": {"model2": "bar"}})
+
+    resolver = DataResolver(dvc, tmp_dir, {"vars": ["test_params.yaml"]})
+    data = {"foreach": "${models}", "do": {"cmd": "echo ${item}"}}
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        "build@model1": {"cmd": "echo foo"},
+        "build@model2": {"cmd": "echo bar"},
+    }
+    assert resolver.context == {"models": {"model1": "foo", "model2": "bar"}}
+    assert resolver.tracked_vars == {
+        "build@model1": {"params.yaml": {"models.model1": "foo"}},
+        "build@model2": {"test_params.yaml": {"models.model2": "bar"}},
+    }
+
+
+def test_mixed_vars_for_foreach_data_2(tmp_dir, dvc):
+    dump_yaml(
+        "params.yaml",
+        {"models": {"model1": {"thresh": 10}, "model2": {"thresh": 15}}},
+    )
+    dump_yaml(
+        "test_params.yaml",
+        {"models": {"model1": {"epochs": 5}, "model2": {"epochs": 10}}},
+    )
+
+    resolver = DataResolver(dvc, tmp_dir, {"vars": ["test_params.yaml"]})
+    data = {
+        "foreach": "${models}",
+        "do": {"cmd": "echo ${item.thresh} ${item.epochs}"},
+    }
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        "build@model1": {"cmd": "echo 10 5"},
+        "build@model2": {"cmd": "echo 15 10"},
+    }
+    assert resolver.context == {
+        "models": {
+            "model1": {"thresh": 10, "epochs": 5},
+            "model2": {"thresh": 15, "epochs": 10},
+        }
+    }
+    assert resolver.tracked_vars == {
+        "build@model1": {
+            "params.yaml": {"models.model1.thresh": 10},
+            "test_params.yaml": {"models.model1.epochs": 5},
+        },
+        "build@model2": {
+            "params.yaml": {"models.model2.thresh": 15},
+            "test_params.yaml": {"models.model2.epochs": 10},
+        },
+    }
+
+
+def test_foreach_with_interpolated_wdir(tmp_dir, dvc):
+    resolver = DataResolver(dvc, tmp_dir / "data", {})
+    foreach_data = ["foo", "bar"]
+    data = {
+        "foreach": foreach_data,
+        "do": {"wdir": "${item}", "cmd": "echo hello"},
+    }
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        # note that the resolver generates `wdir` relative to file's wdir
+        # so, this is just `foo`, not `data/foo`.
+        # figuring out `wdir` is the responsibility of the `load_stage`/`Stage`
+        "build@foo": {"wdir": "foo", "cmd": "echo hello"},
+        "build@bar": {"wdir": "bar", "cmd": "echo hello"},
+    }
+
+    assert not resolver.context
+    assert not any(item for item in resolver.tracked_vars.values())
+
+
+def test_foreach_with_local_vars(tmp_dir, dvc):
+    resolver = DataResolver(dvc, tmp_dir / "data", {})
+    foreach_data = ["foo", "bar"]
+    data = {
+        "foreach": foreach_data,
+        "do": {
+            "vars": [{"foobar": "foobar"}],
+            "cmd": "echo ${item} ${foobar}",
+        },
+    }
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        # note that the resolver generates `wdir` relative to file's wdir
+        # so, this is just `foo`, not `data/foo`.
+        # figuring out `wdir` is the responsibility of the `load_stage`/`Stage`
+        "build@foo": {"cmd": "echo foo foobar"},
+        "build@bar": {"cmd": "echo bar foobar"},
+    }
+    assert not resolver.context
+    assert not any(item for item in resolver.tracked_vars.values())
+
+
+@pytest.mark.parametrize(
+    "local_import",
+    [
+        "test_params.yaml",
+        "test_params.yaml:train",
+        "test_params.yaml:train,prepare",
+    ],
+)
+def test_foreach_with_imported_vars(tmp_dir, dvc, local_import):
+    dump_yaml("params.yaml", {"models": {"model1": {"thresh": "foo"}}})
+    dump_yaml(
+        "test_params.yaml", {"train": {"epochs": 10}, "prepare": {"nums": 25}}
+    )
+    resolver = DataResolver(dvc, tmp_dir, {})
+    foreach_data = ["foo", "bar"]
+    data = {
+        "foreach": foreach_data,
+        "do": {"vars": [local_import], "cmd": "echo ${item} ${train.epochs}"},
+    }
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        # note that the resolver generates `wdir` relative to file's wdir
+        # so, this is just `foo`, not `data/foo`.
+        # figuring out `wdir` is the responsibility of the `load_stage`/`Stage`
+        "build@foo": {"cmd": "echo foo 10"},
+        "build@bar": {"cmd": "echo bar 10"},
+    }
+
+    assert resolver.context == {"models": {"model1": {"thresh": "foo"}}}
+    assert resolver.tracked_vars == {
+        "build@foo": {"test_params.yaml": {"train.epochs": 10}},
+        "build@bar": {"test_params.yaml": {"train.epochs": 10}},
+    }
+
+
+@pytest.mark.parametrize(
+    "local_import", ["params.yaml", "params.yaml:train,prepare"]
+)
+def test_foreach_with_interpolated_wdir_and_local_vars(
+    tmp_dir, dvc, local_import
+):
+    dump_yaml("params.yaml", {"models": {"model1": {"thresh": "foo"}}})
+
+    for i in range(5):
+        build_dir = tmp_dir / ("model-" + str(i))
+        build_dir.mkdir()
+        dump_yaml(
+            build_dir / "params.yaml",
+            {"train": {"epochs": 1 + i}, "prepare": {"nums": 10 * i}},
+        )
+
+    resolver = DataResolver(dvc, tmp_dir, {})
+    data = {
+        "foreach": [0, 1, 2, 3, 4],
+        "do": {
+            "wdir": "model-${item}",
+            "vars": [local_import],
+            "cmd": "echo ${item} ${train.epochs} ${prepare.nums}",
+        },
+    }
+    definition = ForeachDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        # note that the resolver generates `wdir` relative to file's wdir
+        # so, this is just `foo`, not `data/foo`.
+        # figuring out `wdir` is the responsibility of the `load_stage`/`Stage`
+        "build@0": {"wdir": "model-0", "cmd": "echo 0 1 0"},
+        "build@1": {"wdir": "model-1", "cmd": "echo 1 2 10"},
+        "build@2": {"wdir": "model-2", "cmd": "echo 2 3 20"},
+        "build@3": {"wdir": "model-3", "cmd": "echo 3 4 30"},
+        "build@4": {"wdir": "model-4", "cmd": "echo 4 5 40"},
+    }
+
+    assert resolver.context == {"models": {"model1": {"thresh": "foo"}}}
+    assert resolver.tracked_vars == {
+        "build@0": {
+            os.path.join("model-0", "params.yaml"): {
+                "train.epochs": 1,
+                "prepare.nums": 0,
+            }
+        },
+        "build@1": {
+            os.path.join("model-1", "params.yaml"): {
+                "train.epochs": 2,
+                "prepare.nums": 10,
+            }
+        },
+        "build@2": {
+            os.path.join("model-2", "params.yaml"): {
+                "train.epochs": 3,
+                "prepare.nums": 20,
+            }
+        },
+        "build@3": {
+            os.path.join("model-3", "params.yaml"): {
+                "train.epochs": 4,
+                "prepare.nums": 30,
+            }
+        },
+        "build@4": {
+            os.path.join("model-4", "params.yaml"): {
+                "train.epochs": 5,
+                "prepare.nums": 40,
+            }
+        },
+    }
+    assert resolver.context.imports == {
+        str(tmp_dir / DEFAULT_PARAMS_FILE): None
+    }
+
+
+def test_foreach_do_syntax_is_checked_once(tmp_dir, dvc, mocker):
+    do_def = {"cmd": "python script.py --epochs ${item}"}
+    data = {"foreach": [0, 1, 2, 3, 4], "do": do_def}
+    definition = ForeachDefinition(
+        DataResolver(dvc, tmp_dir, {}), Context(), "build", data
+    )
+    mock = mocker.patch("dvc.parsing.check_syntax_errors", return_value=True)
+    definition.resolve_all()
+
+    mock.assert_called_once_with(do_def, "build", "dvc.yaml")
+
+
+def test_foreach_data_is_only_resolved_once(tmp_dir, dvc, mocker):
+    context = Context(models=["foo", "bar", "baz"])
+    data = {"foreach": "${models}", "do": {}}
+    definition = ForeachDefinition(
+        DataResolver(dvc, tmp_dir, {}), context, "build", data
+    )
+    mock = mocker.spy(definition, "_resolve_foreach_data")
+
+    definition.resolve_all()
+
+    mock.assert_called_once_with()

tests/func/parsing/test_interpolated_entry.py

@@ -0,0 +1,262 @@
+from copy import deepcopy
+
+import pytest
+
+from dvc.dependency import _merge_params
+from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver
+from dvc.parsing.context import recurse_not_a_node
+from dvc.path_info import PathInfo
+from dvc.utils.serialize import dump_json, dump_yaml
+
+from . import (
+    CONTEXT_DATA,
+    RESOLVED_DVC_YAML_DATA,
+    TEMPLATED_DVC_YAML_DATA,
+    USED_VARS,
+)
+
+
+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:
+                params[k] = sorted(params[k])
+            if params:
+                stage_d["params"] = params
+    assert d1 == d2
+
+
+def test_simple(tmp_dir, dvc):
+    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, CONTEXT_DATA)
+    resolver = DataResolver(
+        dvc, PathInfo(str(tmp_dir)), deepcopy(TEMPLATED_DVC_YAML_DATA)
+    )
+    assert_stage_equal(resolver.resolve(), deepcopy(RESOLVED_DVC_YAML_DATA))
+    assert resolver.tracked_vars == {
+        "stage1": {DEFAULT_PARAMS_FILE: USED_VARS["stage1"]},
+        "stage2": {DEFAULT_PARAMS_FILE: USED_VARS["stage2"]},
+    }
+
+
+def test_vars_import(tmp_dir, dvc):
+    """
+    Test that different file can be loaded using `vars`
+    instead of default params.yaml.
+    """
+    dump_yaml(tmp_dir / "params2.yaml", CONTEXT_DATA)
+    d = deepcopy(TEMPLATED_DVC_YAML_DATA)
+    d["vars"] = ["params2.yaml"]
+    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
+
+    resolved_data = deepcopy(RESOLVED_DVC_YAML_DATA)
+    assert_stage_equal(resolver.resolve(), resolved_data)
+    assert resolver.tracked_vars == {
+        "stage1": {"params2.yaml": USED_VARS["stage1"]},
+        "stage2": {"params2.yaml": USED_VARS["stage2"]},
+    }
+
+
+def test_vars_and_params_import(tmp_dir, dvc):
+    """
+    Test that vars and params are both merged together for interpolation,
+    whilst tracking the "used" variables from params.
+    """
+    d = {
+        "vars": [DEFAULT_PARAMS_FILE, {"dict": {"foo": "foobar"}}],
+        "stages": {"stage1": {"cmd": "echo ${dict.foo} ${dict.bar}"}},
+    }
+    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, {"dict": {"bar": "bar"}})
+    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
+
+    assert_stage_equal(
+        resolver.resolve(), {"stages": {"stage1": {"cmd": "echo foobar bar"}}}
+    )
+    assert resolver.tracked_vars == {
+        "stage1": {DEFAULT_PARAMS_FILE: {"dict.bar": "bar"}}
+    }
+
+
+def test_stage_with_wdir(tmp_dir, dvc):
+    """
+    Test that params file from wdir are also loaded
+    """
+    d = {
+        "stages": {
+            "stage1": {
+                "cmd": "echo ${dict.foo} ${dict.bar}",
+                "params": ["value1"],
+                "wdir": "data",
+                "vars": [DEFAULT_PARAMS_FILE],
+            },
+        },
+    }
+
+    data_dir = tmp_dir / "data"
+    data_dir.mkdir()
+    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, {"dict": {"bar": "bar"}})
+    dump_json(data_dir / DEFAULT_PARAMS_FILE, {"dict": {"foo": "foo"}})
+    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
+
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                "stage1": {
+                    "cmd": "echo foo bar",
+                    "wdir": "data",
+                    "params": ["value1"],
+                }
+            }
+        },
+    )
+    assert resolver.tracked_vars == {
+        "stage1": {
+            str(PathInfo("data") / DEFAULT_PARAMS_FILE): {"dict.foo": "foo"},
+            DEFAULT_PARAMS_FILE: {"dict.bar": "bar"},
+        }
+    }
+
+
+def test_with_templated_wdir(tmp_dir, dvc):
+    """
+    Test that params from the resolved wdir are still loaded
+    and is used in the interpolation.
+    """
+    d = {
+        "stages": {
+            "stage1": {
+                "cmd": "echo ${dict.foo} ${dict.bar}",
+                "params": ["value1"],
+                "wdir": "${dict.ws}",
+                "vars": [DEFAULT_PARAMS_FILE],
+            },
+        },
+    }
+    dump_yaml(
+        tmp_dir / DEFAULT_PARAMS_FILE, {"dict": {"bar": "bar", "ws": "data"}}
+    )
+    data_dir = tmp_dir / "data"
+    data_dir.mkdir()
+    dump_json(data_dir / DEFAULT_PARAMS_FILE, {"dict": {"foo": "foo"}})
+    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
+
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                "stage1": {
+                    "cmd": "echo foo bar",
+                    "wdir": "data",
+                    "params": ["value1"],
+                }
+            }
+        },
+    )
+    assert resolver.tracked_vars == {
+        "stage1": {
+            str(PathInfo("data") / DEFAULT_PARAMS_FILE): {"dict.foo": "foo"},
+            DEFAULT_PARAMS_FILE: {"dict.bar": "bar", "dict.ws": "data"},
+        }
+    }
+    assert resolver.context.imports == {str(tmp_dir / "params.yaml"): None}
+    assert resolver.context == {"dict": {"bar": "bar", "ws": "data"}}
+
+
+def test_resolve_local_tries_to_load_globally_used_files(tmp_dir, dvc):
+    iterable = {"bar": "bar", "foo": "foo"}
+    dump_json(tmp_dir / "params.json", iterable)
+
+    d = {
+        "vars": ["params.json"],
+        "stages": {
+            "build": {
+                "cmd": "command --value ${bar}",
+                "params": [{"params.json": ["foo"]}],
+                "vars": ["params.json"],
+            },
+        },
+    }
+    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                "build": {
+                    "cmd": "command --value bar",
+                    "params": [{"params.json": ["foo"]}],
+                },
+            }
+        },
+    )
+    assert resolver.tracked_vars == {"build": {"params.json": {"bar": "bar"}}}
+
+
+def test_resolve_local_tries_to_load_globally_used_params_yaml(tmp_dir, dvc):
+    iterable = {"bar": "bar", "foo": "foo"}
+    dump_yaml(tmp_dir / "params.yaml", iterable)
+
+    d = {
+        "stages": {
+            "build": {
+                "cmd": "command --value ${bar}",
+                "params": [{"params.yaml": ["foo"]}],
+                "vars": ["params.yaml"],
+            },
+        },
+    }
+    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
+    assert_stage_equal(
+        resolver.resolve(),
+        {
+            "stages": {
+                "build": {
+                    "cmd": "command --value bar",
+                    "params": [{"params.yaml": ["foo"]}],
+                },
+            }
+        },
+    )
+    assert resolver.tracked_vars == {"build": {"params.yaml": {"bar": "bar"}}}
+
+
+def test_vars_relpath_overwrite(tmp_dir, dvc):
+    iterable = {"bar": "bar", "foo": "foo"}
+    dump_yaml(tmp_dir / "params.yaml", iterable)
+    d = {
+        "vars": ["params.yaml"],
+        "stages": {
+            "build": {
+                "wdir": "data",
+                "cmd": "echo ${bar}",
+                "vars": ["../params.yaml"],
+            }
+        },
+    }
+    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
+    resolver.resolve()
+    assert resolver.context.imports == {str(tmp_dir / "params.yaml"): None}
+
+
+@pytest.mark.parametrize("local", [True, False])
+@pytest.mark.parametrize(
+    "vars_",
+    [
+        ["test_params.yaml:bar", "test_params.yaml:foo"],
+        ["test_params.yaml:foo,bar"],
+        ["test_params.yaml"],
+        ["test_params.yaml", "test_params.yaml"],
+    ],
+)
+def test_vars_load_partial(tmp_dir, dvc, local, vars_):
+    iterable = {"bar": "bar", "foo": "foo"}
+    dump_yaml(tmp_dir / "test_params.yaml", iterable)
+    d = {"stages": {"build": {"cmd": "echo ${bar}"}}}
+    if local:
+        d["stages"]["build"]["vars"] = vars_
+    else:
+        d["vars"] = vars_
+    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
+    resolver.resolve()

tests/func/parsing/test_resolver.py

@@ -0,0 +1,177 @@
+from copy import deepcopy
+
+import pytest
+
+from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver, ResolveError
+from dvc.parsing.context import Context
+from dvc.utils.serialize import dump_json, dump_yaml
+
+from . import CONTEXT_DATA, RESOLVED_DVC_YAML_DATA, TEMPLATED_DVC_YAML_DATA
+
+DATA = {"models": {"bar": "bar", "foo": "foo"}}
+
+
+def test_resolver(tmp_dir, dvc):
+    resolver = DataResolver(dvc, tmp_dir, TEMPLATED_DVC_YAML_DATA)
+    resolver.context.merge_update(Context(CONTEXT_DATA))
+    assert resolver.resolve() == RESOLVED_DVC_YAML_DATA
+
+
+def test_default_params_file_not_exist(tmp_dir, dvc):
+    d = {"vars": [DATA["models"]]}
+    resolver = DataResolver(dvc, tmp_dir, d)
+    assert resolver.context == d["vars"][0]
+
+
+def test_no_params_yaml_and_vars(tmp_dir, dvc):
+    resolver = DataResolver(dvc, tmp_dir, {})
+    assert not resolver.context
+
+
+def test_local_vars(tmp_dir, dvc):
+    resolver = DataResolver(
+        dvc, tmp_dir, {"vars": [{"foo": "bar", "bar": "foo"}]}
+    )
+    assert resolver.context == {"foo": "bar", "bar": "foo"}
+
+
+@pytest.mark.parametrize("vars_", ["${file}_params.yaml", {"foo": "${foo}"}])
+def test_vars_interpolation_errors(tmp_dir, dvc, vars_):
+    with pytest.raises(ResolveError) as exc_info:
+        DataResolver(dvc, tmp_dir, {"vars": [vars_, {"bar": "foo"}]})
+    assert (
+        str(exc_info.value)
+        == "failed to parse 'vars' in 'dvc.yaml': interpolating is not allowed"
+    )
+
+
+@pytest.mark.parametrize(
+    "vars_", [{}, {"vars": []}, {"vars": [DEFAULT_PARAMS_FILE]}]
+)
+def test_default_params_file(tmp_dir, dvc, vars_):
+    dump_yaml(DEFAULT_PARAMS_FILE, DATA)
+    resolver = DataResolver(dvc, tmp_dir, vars_)
+    assert resolver.context == DATA
+
+
+def test_load_vars_from_file(tmp_dir, dvc):
+    dump_yaml(DEFAULT_PARAMS_FILE, DATA)
+
+    datasets = {"datasets": ["foo", "bar"]}
+    dump_json("params.json", datasets)
+    d = {"vars": [DEFAULT_PARAMS_FILE, "params.json"]}
+    resolver = DataResolver(dvc, tmp_dir, d)
+
+    expected = deepcopy(DATA)
+    expected.update(datasets)
+    assert resolver.context == expected
+
+
+def test_partial_vars_doesnot_exist(tmp_dir, dvc):
+    dump_yaml("test_params.yaml", {"sub1": "sub1"})
+
+    with pytest.raises(ResolveError) as exc_info:
+        DataResolver(dvc, tmp_dir, {"vars": ["test_params.yaml:sub2"]})
+
+    assert (
+        str(exc_info.value) == "failed to parse 'vars' in 'dvc.yaml': "
+        "could not find 'sub2' in 'test_params.yaml'"
+    )
+
+
+def test_global_overwrite_error_on_imports(tmp_dir, dvc):
+    dump_yaml(DEFAULT_PARAMS_FILE, DATA)
+    dump_json("params.json", DATA)
+
+    d = {"vars": [DEFAULT_PARAMS_FILE, "params.json"]}
+    with pytest.raises(ResolveError) as exc_info:
+        DataResolver(dvc, tmp_dir, d)
+
+    assert str(exc_info.value) == (
+        "failed to parse 'vars' in 'dvc.yaml':\n"
+        "cannot redefine 'models.bar' from 'params.json' "
+        "as it already exists in 'params.yaml'"
+    )
+
+
+def test_global_overwrite_vars(tmp_dir, dvc):
+    dump_yaml(DEFAULT_PARAMS_FILE, DATA)
+    d = {"vars": [DATA]}
+
+    with pytest.raises(ResolveError) as exc_info:
+        DataResolver(dvc, tmp_dir, d)
+
+    assert str(exc_info.value) == (
+        "failed to parse 'vars' in 'dvc.yaml':\n"
+        "cannot redefine 'models.bar' from 'vars[0]' "
+        "as it already exists in 'params.yaml'"
+    )
+
+
+def test_local_declared_vars_overwrite(tmp_dir, dvc):
+    dump_yaml(DEFAULT_PARAMS_FILE, DATA)
+
+    d = {"vars": [DATA["models"], DATA["models"]]}
+    with pytest.raises(ResolveError) as exc_info:
+        DataResolver(dvc, tmp_dir, d)
+
+    assert str(exc_info.value) == (
+        "failed to parse 'vars' in 'dvc.yaml':\n"
+        "cannot redefine 'bar' from 'vars[1]' "
+        "as it already exists in 'vars[0]'"
+    )
+
+
+def test_specified_params_file_not_exist(tmp_dir, dvc):
+    d = {"vars": ["not_existing_params.yaml"]}
+    with pytest.raises(ResolveError) as exc_info:
+        DataResolver(dvc, tmp_dir, d)
+
+    assert str(exc_info.value) == (
+        "failed to parse 'vars' in 'dvc.yaml': "
+        "'not_existing_params.yaml' does not exist"
+    )
+
+
+@pytest.mark.parametrize("local", [True, False])
+@pytest.mark.parametrize(
+    "vars_",
+    [
+        ["test_params.yaml", "test_params.yaml:sub1"],
+        ["test_params.yaml:sub1", "test_params.yaml"],
+        ["test_params.yaml:sub1", "test_params.yaml:sub1,sub2"],
+    ],
+)
+def test_vars_already_loaded_message(tmp_dir, dvc, local, vars_):
+    d = {"stages": {"build": {"cmd": "echo ${sub1} ${sub2}"}}}
+    dump_yaml("test_params.yaml", {"sub1": "sub1", "sub2": "sub2"})
+    if not local:
+        d["vars"] = vars_
+    else:
+        d["stages"]["build"]["vars"] = vars_
+
+    with pytest.raises(ResolveError) as exc_info:
+        resolver = DataResolver(dvc, tmp_dir, d)
+        resolver.resolve()
+
+    assert "partially" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+    "vars_, loc", [(DATA, "build.vars[0]"), ("params.json", "params.json")],
+)
+def test_local_overwrite_error(tmp_dir, dvc, vars_, loc):
+    dump_yaml(DEFAULT_PARAMS_FILE, DATA)
+    dump_json("params.json", DATA)
+
+    d = {"stages": {"build": {"cmd": "echo ${models.foo}", "vars": [vars_]}}}
+
+    resolver = DataResolver(dvc, tmp_dir, d)
+    with pytest.raises(ResolveError) as exc_info:
+        resolver.resolve()
+
+    assert str(exc_info.value) == (
+        "failed to parse stage 'build' in 'dvc.yaml':\n"
+        f"cannot redefine 'models.bar' from '{loc}' "
+        "as it already exists in 'params.yaml'"
+    )

tests/func/test_stage_resolver.py

@@ -1,536 +0,0 @@
-from copy import deepcopy
-from math import pi
-
-import pytest
-
-from dvc.dependency import _merge_params
-from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver, ResolveError
-from dvc.parsing.context import Node
-from dvc.path_info import PathInfo
-from dvc.utils.serialize import dump_json, dump_yaml
-
-TEMPLATED_DVC_YAML_DATA = {
-    "stages": {
-        "stage1": {
-            "cmd": "python script.py ${dict.foo} --out ${dict.bar}",
-            "outs": ["${dict.bar}"],
-            "deps": ["${dict.foo}"],
-            "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"],
-            "frozen": True,
-        },
-        "stage2": {"cmd": "echo foo bar"},
-    }
-}
-
-USED_VARS = {
-    "stage1": {"dict.foo": "foo", "dict.bar": "bar", "freeze": True},
-    "stage2": {"dict.foo": "foo", "dict.bar": "bar"},
-}
-
-
-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:
-                params[k] = sorted(params[k])
-            if params:
-                stage_d["params"] = params
-    assert d1 == d2
-
-
-def test_simple(tmp_dir, dvc):
-    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, CONTEXT_DATA)
-    resolver = DataResolver(
-        dvc, PathInfo(str(tmp_dir)), deepcopy(TEMPLATED_DVC_YAML_DATA)
-    )
-    assert_stage_equal(resolver.resolve(), deepcopy(RESOLVED_DVC_YAML_DATA))
-    assert resolver.tracked_vars == {
-        "stage1": {DEFAULT_PARAMS_FILE: USED_VARS["stage1"]},
-        "stage2": {DEFAULT_PARAMS_FILE: USED_VARS["stage2"]},
-    }
-
-
-def test_vars(tmp_dir, dvc):
-    d = deepcopy(TEMPLATED_DVC_YAML_DATA)
-    d["vars"] = [CONTEXT_DATA]
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    resolved_data = deepcopy(RESOLVED_DVC_YAML_DATA)
-
-    assert_stage_equal(resolver.resolve(), resolved_data)
-    assert not any(resolver.tracked_vars.values())
-
-
-def test_no_params_yaml_and_vars(tmp_dir, dvc):
-    resolver = DataResolver(
-        dvc, PathInfo(str(tmp_dir)), deepcopy(TEMPLATED_DVC_YAML_DATA)
-    )
-    with pytest.raises(ResolveError):
-        resolver.resolve()
-
-
-def test_vars_import(tmp_dir, dvc):
-    """
-    Test that different file can be loaded using `vars`
-    instead of default params.yaml.
-    """
-    dump_yaml(tmp_dir / "params2.yaml", CONTEXT_DATA)
-    d = deepcopy(TEMPLATED_DVC_YAML_DATA)
-    d["vars"] = ["params2.yaml"]
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-
-    resolved_data = deepcopy(RESOLVED_DVC_YAML_DATA)
-    assert_stage_equal(resolver.resolve(), resolved_data)
-    assert resolver.tracked_vars == {
-        "stage1": {"params2.yaml": USED_VARS["stage1"]},
-        "stage2": {"params2.yaml": USED_VARS["stage2"]},
-    }
-
-
-def test_vars_and_params_import(tmp_dir, dvc):
-    """
-    Test that vars and params are both merged together for interpolation,
-    whilst tracking the "used" variables from params.
-    """
-    d = {
-        "vars": [DEFAULT_PARAMS_FILE, {"dict": {"foo": "foobar"}}],
-        "stages": {"stage1": {"cmd": "echo ${dict.foo} ${dict.bar}"}},
-    }
-    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, {"dict": {"bar": "bar"}})
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-
-    assert_stage_equal(
-        resolver.resolve(), {"stages": {"stage1": {"cmd": "echo foobar bar"}}}
-    )
-    assert resolver.tracked_vars == {
-        "stage1": {DEFAULT_PARAMS_FILE: {"dict.bar": "bar"}}
-    }
-
-
-def test_with_params_section(tmp_dir, dvc):
-    """Test that params section is also loaded for interpolation"""
-    d = {
-        "vars": [DEFAULT_PARAMS_FILE, {"dict": {"foo": "foo"}}],
-        "stages": {
-            "stage1": {
-                "cmd": "echo ${dict.foo} ${dict.bar} ${dict.foobar}",
-                "params": [{"params.json": ["value1"]}],
-                "vars": ["params.json"],
-            },
-        },
-    }
-    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, {"dict": {"bar": "bar"}})
-    dump_json(tmp_dir / "params.json", {"dict": {"foobar": "foobar"}})
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "stage1": {
-                    "cmd": "echo foo bar foobar",
-                    "params": [{"params.json": ["value1"]}],
-                }
-            }
-        },
-    )
-    assert resolver.tracked_vars == {
-        "stage1": {
-            DEFAULT_PARAMS_FILE: {"dict.bar": "bar"},
-            "params.json": {"dict.foobar": "foobar"},
-        }
-    }
-
-
-def test_stage_with_wdir(tmp_dir, dvc):
-    """
-    Test that params file from wdir are also loaded
-    """
-    d = {
-        "stages": {
-            "stage1": {
-                "cmd": "echo ${dict.foo} ${dict.bar}",
-                "params": ["value1"],
-                "wdir": "data",
-                "vars": [DEFAULT_PARAMS_FILE],
-            },
-        },
-    }
-
-    data_dir = tmp_dir / "data"
-    data_dir.mkdir()
-    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, {"dict": {"bar": "bar"}})
-    dump_json(data_dir / DEFAULT_PARAMS_FILE, {"dict": {"foo": "foo"}})
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "stage1": {
-                    "cmd": "echo foo bar",
-                    "wdir": "data",
-                    "params": ["value1"],
-                }
-            }
-        },
-    )
-    assert resolver.tracked_vars == {
-        "stage1": {
-            str(PathInfo("data") / DEFAULT_PARAMS_FILE): {"dict.foo": "foo"},
-            DEFAULT_PARAMS_FILE: {"dict.bar": "bar"},
-        }
-    }
-
-
-def test_with_templated_wdir(tmp_dir, dvc):
-    """
-    Test that params from the resolved wdir are still loaded
-    and is used in the interpolation.
-    """
-    d = {
-        "stages": {
-            "stage1": {
-                "cmd": "echo ${dict.foo} ${dict.bar}",
-                "params": ["value1"],
-                "wdir": "${dict.ws}",
-                "vars": [DEFAULT_PARAMS_FILE],
-            },
-        },
-    }
-    dump_yaml(
-        tmp_dir / DEFAULT_PARAMS_FILE, {"dict": {"bar": "bar", "ws": "data"}}
-    )
-    data_dir = tmp_dir / "data"
-    data_dir.mkdir()
-    dump_json(data_dir / DEFAULT_PARAMS_FILE, {"dict": {"foo": "foo"}})
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "stage1": {
-                    "cmd": "echo foo bar",
-                    "wdir": "data",
-                    "params": ["value1"],
-                }
-            }
-        },
-    )
-    assert resolver.tracked_vars == {
-        "stage1": {
-            str(PathInfo("data") / DEFAULT_PARAMS_FILE): {"dict.foo": "foo"},
-            DEFAULT_PARAMS_FILE: {"dict.bar": "bar", "dict.ws": "data"},
-        }
-    }
-
-
-def test_simple_foreach_loop(tmp_dir, dvc):
-    iterable = ["foo", "bar", "baz"]
-    d = {
-        "stages": {
-            "build": {
-                "foreach": iterable,
-                "do": {"cmd": "python script.py ${item}"},
-            }
-        }
-    }
-
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    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):
-    iterable = {"models": {"us": {"thresh": 10}, "gb": {"thresh": 15}}}
-    d = {
-        "stages": {
-            "build": {
-                "foreach": iterable["models"],
-                "do": {"cmd": "python script.py ${item.thresh}"},
-            }
-        }
-    }
-
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    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):
-    params = {"models": {"us": {"thresh": 10}}}
-    vars_ = [{"models": {"gb": {"thresh": 15}}}]
-    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, params)
-    d = {
-        "vars": vars_,
-        "stages": {
-            "build": {
-                "foreach": "${models}",
-                "do": {"cmd": "python script.py --thresh ${item.thresh}"},
-            }
-        },
-    }
-
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "build@gb": {"cmd": "python script.py --thresh 15"},
-                "build@us": {"cmd": "python script.py --thresh 10"},
-            }
-        },
-    )
-    assert resolver.tracked_vars == {
-        "build@gb": {},
-        "build@us": {DEFAULT_PARAMS_FILE: {"models.us.thresh": 10}},
-    }
-
-
-@pytest.mark.parametrize(
-    "value", ["value", "To set or not to set", 3, pi, True, False, None]
-)
-def test_set(tmp_dir, dvc, value):
-    d = {
-        "stages": {
-            "build": {
-                "set": {"item": value},
-                "cmd": "python script.py --thresh ${item}",
-                "always_changed": "${item}",
-            }
-        }
-    }
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    if isinstance(value, bool):
-        stringified_value = "true" if value else "false"
-    else:
-        stringified_value = str(value)
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "build": {
-                    "cmd": f"python script.py --thresh {stringified_value}",
-                    "always_changed": value,
-                }
-            }
-        },
-    )
-
-
-@pytest.mark.parametrize(
-    "coll", [["foo", "bar", "baz"], {"foo": "foo", "bar": "bar"}]
-)
-def test_coll(tmp_dir, dvc, coll):
-    d = {
-        "stages": {
-            "build": {
-                "set": {"item": coll, "thresh": 10},
-                "cmd": "python script.py --thresh ${thresh}",
-                "outs": "${item}",
-            }
-        }
-    }
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "build": {"cmd": "python script.py --thresh 10", "outs": coll}
-            }
-        },
-    )
-
-
-def test_set_with_foreach(tmp_dir, dvc):
-    items = ["foo", "bar", "baz"]
-    d = {
-        "stages": {
-            "build": {
-                "set": {"items": items},
-                "foreach": "${items}",
-                "do": {"cmd": "command --value ${item}"},
-            }
-        }
-    }
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    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):
-    iterable = {"models": {"us": {"thresh": 10}, "gb": {"thresh": 15}}}
-    dump_json(tmp_dir / "params.json", iterable)
-
-    d = {
-        "vars": ["params.json"],
-        "stages": {
-            "build": {
-                "set": {"data": "${models}"},
-                "foreach": "${data}",
-                "do": {
-                    "set": {"thresh": "${item.thresh}"},
-                    "cmd": "command --value ${thresh}",
-                },
-            }
-        },
-    }
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "build@us": {"cmd": "command --value 10"},
-                "build@gb": {"cmd": "command --value 15"},
-            }
-        },
-    )
-    assert resolver.tracked_vars == {
-        "build@us": {"params.json": {"models.us.thresh": 10}},
-        "build@gb": {"params.json": {"models.gb.thresh": 15}},
-    }
-
-
-def test_resolve_local_tries_to_load_globally_used_files(tmp_dir, dvc):
-    iterable = {"bar": "bar", "foo": "foo"}
-    dump_json(tmp_dir / "params.json", iterable)
-
-    d = {
-        "vars": ["params.json"],
-        "stages": {
-            "build": {
-                "cmd": "command --value ${bar}",
-                "params": [{"params.json": ["foo"]}],
-                "vars": ["params.json"],
-            },
-        },
-    }
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "build": {
-                    "cmd": "command --value bar",
-                    "params": [{"params.json": ["foo"]}],
-                },
-            }
-        },
-    )
-    assert resolver.tracked_vars == {"build": {"params.json": {"bar": "bar"}}}
-
-
-def test_resolve_local_tries_to_load_globally_used_params_yaml(tmp_dir, dvc):
-    iterable = {"bar": "bar", "foo": "foo"}
-    dump_yaml(tmp_dir / "params.yaml", iterable)
-
-    d = {
-        "stages": {
-            "build": {
-                "cmd": "command --value ${bar}",
-                "params": [{"params.yaml": ["foo"]}],
-                "vars": ["params.yaml"],
-            },
-        },
-    }
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    assert_stage_equal(
-        resolver.resolve(),
-        {
-            "stages": {
-                "build": {
-                    "cmd": "command --value bar",
-                    "params": [{"params.yaml": ["foo"]}],
-                },
-            }
-        },
-    )
-    assert resolver.tracked_vars == {"build": {"params.yaml": {"bar": "bar"}}}
-
-
-def test_vars_relpath_overwrite(tmp_dir, dvc):
-    iterable = {"bar": "bar", "foo": "foo"}
-    dump_yaml(tmp_dir / "params.yaml", iterable)
-    d = {
-        "vars": ["params.yaml"],
-        "stages": {
-            "build": {
-                "wdir": "data",
-                "cmd": "echo ${bar}",
-                "vars": ["../params.yaml"],
-            }
-        },
-    }
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    resolver.resolve()
-
-
-@pytest.mark.parametrize("local", [True, False])
-@pytest.mark.parametrize(
-    "vars_",
-    [
-        ["test_params.yaml:bar", "test_params.yaml:foo"],
-        ["test_params.yaml:foo,bar"],
-        ["test_params.yaml"],
-        ["test_params.yaml", "test_params.yaml"],
-    ],
-)
-def test_vars_load_partial(tmp_dir, dvc, local, vars_):
-    iterable = {"bar": "bar", "foo": "foo"}
-    dump_yaml(tmp_dir / "test_params.yaml", iterable)
-    d = {"stages": {"build": {"cmd": "echo ${bar}"}}}
-    if local:
-        d["stages"]["build"]["vars"] = vars_
-    else:
-        d["vars"] = vars_
-    resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d)
-    resolver.resolve()

tests/func/test_stage_resolver_error_msg.py

@@ -1,384 +0,0 @@
-import logging
-import re
-
-import pytest
-
-from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver, ResolveError
-from dvc.utils.fs import remove
-from dvc.utils.serialize import dump_json, dump_yaml
-
-DATA = {"models": {"bar": "bar", "foo": "foo"}}
-
-
-def escape_ansi(line):
-    ansi_escape = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]")
-    return ansi_escape.sub("", line)
-
-
-@pytest.fixture
-def repo(tmp_dir, dvc):
-    dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, DATA)
-    dump_json(tmp_dir / "params.json", DATA)
-    yield dvc
-
-
-def test_global_overwrite_error_on_imports(tmp_dir, repo):
-    d = {"vars": [DEFAULT_PARAMS_FILE, "params.json"]}
-    with pytest.raises(ResolveError) as exc_info:
-        DataResolver(repo, tmp_dir, d)
-    assert str(exc_info.value) == (
-        "failed to parse 'vars' in 'dvc.yaml':\n"
-        "cannot redefine 'models.bar' from 'params.json' "
-        "as it already exists in 'params.yaml'"
-    )
-
-
-def test_global_overwrite_vars(tmp_dir, repo):
-    d = {"vars": [DATA]}
-    with pytest.raises(ResolveError) as exc_info:
-        DataResolver(repo, tmp_dir, d)
-    assert str(exc_info.value) == (
-        "failed to parse 'vars' in 'dvc.yaml':\n"
-        "cannot redefine 'models.bar' from 'vars[0]' "
-        "as it already exists in 'params.yaml'"
-    )
-
-
-def test_local_declared_vars_overwrite(tmp_dir, repo):
-    d = {"vars": [DATA["models"], DATA["models"]]}
-    with pytest.raises(ResolveError) as exc_info:
-        DataResolver(repo, tmp_dir, d)
-    assert str(exc_info.value) == (
-        "failed to parse 'vars' in 'dvc.yaml':\n"
-        "cannot redefine 'bar' from 'vars[1]' "
-        "as it already exists in 'vars[0]'"
-    )
-
-
-def test_default_params_file_not_exist(tmp_dir, repo):
-    remove(tmp_dir / DEFAULT_PARAMS_FILE)
-    d = {"vars": [DATA["models"]]}
-    DataResolver(repo, tmp_dir, d)
-
-
-def test_specified_params_file_not_exist(tmp_dir, repo):
-    d = {"vars": ["not_existing_params.yaml"]}
-    with pytest.raises(ResolveError) as exc_info:
-        DataResolver(repo, tmp_dir, d)
-    assert str(exc_info.value) == (
-        "failed to parse 'vars' in 'dvc.yaml': "
-        "'not_existing_params.yaml' does not exist"
-    )
-
-
-def test_failed_to_interpolate(tmp_dir, repo):
-    d = {
-        "vars": [DEFAULT_PARAMS_FILE],
-        "stages": {"build": {"cmd": "echo ${models.foo.}"}},
-    }
-    with pytest.raises(ResolveError) as exc_info:
-        resolver = DataResolver(repo, tmp_dir, d)
-        resolver.resolve()
-    assert escape_ansi(str(exc_info.value)) == (
-        "failed to parse 'stages.build.cmd' in 'dvc.yaml':\n"
-        "${models.foo.}\n"
-        "            ^\n"
-        "ParseException: Expected end of text, found '.'"
-        "  (at char 12), (line:1, col:13)"
-    )
-
-
-@pytest.mark.parametrize(
-    "vars_, loc", [(DATA, "build.vars[0]"), ("params.json", "params.json")],
-)
-def test_local_overwrite_error(tmp_dir, repo, vars_, loc):
-    d = {"stages": {"build": {"cmd": "echo ${models.foo}", "vars": [vars_]}}}
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert str(exc_info.value) == (
-        "failed to parse stage 'build' in 'dvc.yaml':\n"
-        f"cannot redefine 'models.bar' from '{loc}' "
-        "as it already exists in 'params.yaml'"
-    )
-
-
-def test_local_vars_params_file_not_exist(tmp_dir, repo):
-    d = {
-        "stages": {
-            "build": {
-                "cmd": "echo ${models.foo}",
-                "vars": ["not_existing_params.yaml"],
-            }
-        }
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-        assert str(exc_info.value) == (
-            "failed to parse stage 'build' in 'dvc.yaml': "
-            "'not_existing_params.yaml' does not exist"
-        )
-
-
-def test_specified_key_does_not_exist(tmp_dir, repo):
-    d = {"stages": {"build": {"cmd": "echo ${models.foobar}"}}}
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert str(exc_info.value) == (
-        "failed to parse 'stages.build.cmd' in 'dvc.yaml': "
-        "Could not find 'models.foobar'"
-    )
-
-
-@pytest.mark.parametrize(
-    "local_vars, expected_text",
-    [
-        ({"item": "item"}, "item is"),
-        ({"key": "key"}, "key is"),
-        ({"item": "item", "key": "key"}, "item and key are"),
-    ],
-)
-def test_foreach_item_key_already_exists(
-    tmp_dir, repo, caplog, local_vars, expected_text
-):
-    d = {
-        "vars": [local_vars],
-        "stages": {"build": {"foreach": "${models}", "do": {}}},
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-
-    caplog.clear()
-    with caplog.at_level(logging.WARNING, logger="dvc.parsing"):
-        resolver.resolve()
-    assert expected_text in caplog.text
-    assert "generated from 'build'" in caplog.text
-
-
-def test_foreach_interpolation_key_does_not_exist(tmp_dir, repo):
-    d = {
-        "stages": {"build": {"foreach": "${modelss}", "do": {}}},
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert str(exc_info.value) == (
-        "failed to parse 'stages.build.foreach' in 'dvc.yaml': "
-        "Could not find 'modelss'"
-    )
-
-
-def test_foreach_interpolation_key_error(tmp_dir, repo):
-    d = {
-        "stages": {"build": {"foreach": "${models[123}", "do": {}}},
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert escape_ansi(str(exc_info.value)) == (
-        "failed to parse 'stages.build.foreach' in 'dvc.yaml':\n"
-        "${models[123}\n"
-        "        ^\n"
-        "ParseException: Expected end of text, found '['"
-        "  (at char 8), (line:1, col:9)"
-    )
-
-
-@pytest.mark.parametrize(
-    "syn, msg",
-    [
-        ("${ite}", " Could not find 'ite'"),
-        (
-            "${ite.}",
-            "\n${ite.}\n"
-            "     ^\n"
-            "ParseException: Expected end of text, found '.'"
-            "  (at char 5), (line:1, col:6)",
-        ),
-    ],
-)
-def test_foreach_generated_errors(tmp_dir, repo, syn, msg):
-    d = {
-        "stages": {
-            "build": {"foreach": "${models}", "do": {"cmd": f"echo {syn}"}}
-        },
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert escape_ansi(str(exc_info.value)) == (
-        "failed to parse '[email protected]' in 'dvc.yaml':" + msg
-    )
-
-
-@pytest.mark.parametrize(
-    "vars_, msg",
-    [
-        (["params-does-not-exist"], " 'params-does-not-exist' does not exist"),
-        (
-            [DATA],
-            "\ncannot redefine 'models.bar' from '[email protected][0]' "
-            f"as it already exists in '{DEFAULT_PARAMS_FILE}'",
-        ),
-    ],
-)
-def test_foreach_generated_local_vars_error(tmp_dir, repo, vars_, msg):
-    d = {
-        "stages": {
-            "build": {
-                "foreach": "${models}",
-                "do": {"vars": vars_, "cmd": "echo ${item}"},
-            }
-        },
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert str(exc_info.value) == (
-        "failed to parse stage 'build@bar' (gen. from 'build') in 'dvc.yaml':"
-        + msg
-    )
-
-
-@pytest.mark.parametrize(
-    "wdir, msg",
-    [
-        ("${ite}", " Could not find 'ite'"),
-        (
-            "${item.}",
-            "\n${item.}\n"
-            "      ^\n"
-            "ParseException: Expected end of text, found '.'"
-            "  (at char 6), (line:1, col:7)",
-        ),
-    ],
-)
-def test_foreach_wdir_interpolation_issues(tmp_dir, repo, wdir, msg):
-    d = {
-        "stages": {
-            "build": {
-                "foreach": "${models}",
-                "do": {"wdir": wdir, "cmd": "echo ${item}"},
-            }
-        },
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert escape_ansi(str(exc_info.value)) == (
-        "failed to parse '[email protected]' in 'dvc.yaml':" + msg
-    )
-
-
-@pytest.mark.parametrize("foreach", ["${models.bar}", "foo"])
-def test_foreach_expects_dict_or_list(tmp_dir, repo, foreach):
-    d = {
-        "stages": {"build": {"foreach": foreach, "do": {}}},
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert str(exc_info.value) == (
-        "failed to resolve 'stages.build.foreach' in 'dvc.yaml': "
-        "expected list/dictionary, got str"
-    )
-
-
-@pytest.mark.parametrize(
-    "wdir, expected_msg",
-    [
-        ("${models[foobar]}", " Could not find 'models.foobar'"),
-        (
-            "${models.foo]}",
-            "\n${models.foo]}\n"
-            "            ^\n"
-            "ParseException: Expected end of text, found ']'"
-            "  (at char 12), (line:1, col:13)",
-        ),
-    ],
-)
-def test_wdir_failed_to_interpolate(tmp_dir, repo, wdir, expected_msg):
-    d = {"stages": {"build": {"wdir": wdir, "cmd": "echo ${models.bar}"}}}
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert escape_ansi(str(exc_info.value)) == (
-        "failed to parse 'stages.build.wdir' in 'dvc.yaml':" + expected_msg
-    )
-
-
-@pytest.mark.parametrize(
-    "foreach, redefine",
-    [
-        ("${models}", "item"),
-        ("${models}", "key"),
-        (list(DATA["models"].keys()), "item"),
-    ],
-)
-def test_item_key_in_generated_stage_vars(tmp_dir, repo, foreach, redefine):
-    d = {
-        "stages": {
-            "build": {
-                "foreach": foreach,
-                "do": {"cmd": "echo ${item}", "vars": [{redefine: "value"}]},
-            }
-        },
-    }
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert str(exc_info.value) == (
-        f"attempted to redefine '{redefine}' in stage 'build@bar'"
-        " generated through 'foreach'"
-    )
-
-
-def test_interpolate_non_string(tmp_dir, repo):
-    d = {"stages": {"build": {"cmd": "echo ${models}"}}}
-    resolver = DataResolver(repo, tmp_dir, d)
-    with pytest.raises(ResolveError) as exc_info:
-        resolver.resolve()
-    assert str(exc_info.value) == (
-        "failed to parse 'stages.build.cmd' in 'dvc.yaml':\n"
-        "Cannot interpolate data of type 'dict'"
-    )
-
-
-@pytest.mark.parametrize("local", [True, False])
-@pytest.mark.parametrize(
-    "vars_",
-    [
-        ["test_params.yaml", "test_params.yaml:sub1"],
-        ["test_params.yaml:sub1", "test_params.yaml"],
-        ["test_params.yaml:sub1", "test_params.yaml:sub1,sub2"],
-    ],
-)
-def test_vars_already_loaded_message(tmp_dir, repo, local, vars_):
-    d = {"stages": {"build": {"cmd": "echo ${sub1} ${sub2}"}}}
-    dump_yaml("test_params.yaml", {"sub1": "sub1", "sub2": "sub2"})
-    if not local:
-        d["vars"] = vars_
-    else:
-        d["stages"]["build"]["vars"] = vars_
-
-    with pytest.raises(ResolveError) as exc_info:
-        resolver = DataResolver(repo, tmp_dir, d)
-        resolver.resolve()
-
-    assert "partially" in str(exc_info.value)
-
-
-@pytest.mark.parametrize("local", [False, True])
-def test_partial_vars_doesnot_exist(tmp_dir, repo, local):
-    d = {"stages": {"build": {"cmd": "echo ${sub1} ${sub2}"}}}
-    dump_yaml("test_params.yaml", {"sub1": "sub1", "sub2": "sub2"})
-    vars_ = ["test_params.yaml:sub3"]
-    if not local:
-        d["vars"] = vars_
-    else:
-        d["stages"]["build"]["vars"] = vars_
-
-    with pytest.raises(ResolveError):
-        resolver = DataResolver(repo, tmp_dir, d)
-        resolver.resolve()

tests/unit/test_context.py

@@ -12,11 +12,11 @@ from dvc.parsing.context import (
     MergeError,
     ParamsFileNotFound,
     Value,
+    recurse_not_a_node,
 )
 from dvc.tree.local import LocalTree
 from dvc.utils import relpath
 from dvc.utils.serialize import dump_yaml
-from tests.func.test_stage_resolver import recurse_not_a_node
 
 
 def test_context():
@@ -262,22 +262,25 @@ def test_load_from(mocker):
 
 def test_clone():
     d = {
-        "lst": [
-            {"foo0": "foo0", "bar0": "bar0"},
-            {"foo1": "foo1", "bar1": "bar1"},
-        ]
+        "dct": {
+            "foo0": "foo0",
+            "bar0": "bar0",
+            "foo1": "foo1",
+            "bar1": "bar1",
+        },
+        "lst": [1, 2, 3],
     }
     c1 = Context(d)
     c2 = Context.clone(c1)
 
-    c2["lst"][0]["foo0"] = "foo"
-    del c2["lst"][1]["foo1"]
+    c2["dct"]["foo0"] = "foo"
+    del c2["dct"]["foo1"]
 
     assert c1 != c2
     assert c1 == Context(d)
-    assert c2.select("lst.0.foo0") == Value("foo")
+    assert c2.select("lst.0") == Value(1)
     with pytest.raises(KeyNotInContext):
-        c2.select("lst.1.foo1")
+        c2.select("lst.1.not_existing_key")
 
 
 def test_track(tmp_dir):
@@ -294,29 +297,29 @@ def test_track(tmp_dir):
 
     context = Context.load_from(tree, path)
 
-    def key_tracked(key):
-        assert len(context.tracked) == 1
-        return key in context.tracked[relpath(path)]
+    def key_tracked(d, key):
+        assert len(d) == 1
+        return key in d[relpath(path)]
 
-    with context.track():
+    with context.track() as tracked:
         context.select("lst")
-        assert key_tracked("lst")
+        assert key_tracked(tracked, "lst")
 
         context.select("dct")
-        assert not key_tracked("dct")
+        assert not key_tracked(tracked, "dct")
 
         context.select("dct.foo")
-        assert key_tracked("dct.foo")
+        assert key_tracked(tracked, "dct.foo")
 
         # Currently, it's unable to track dictionaries, as it can be merged
         # from multiple sources.
         context.select("lst.0")
-        assert not key_tracked("lst.0")
+        assert not key_tracked(tracked, "lst.0")
 
         # FIXME: either support tracking list values in ParamsDependency
         # or, prevent this from being tracked.
         context.select("lst.0.foo0")
-        assert key_tracked("lst.0.foo0")
+        assert key_tracked(tracked, "lst.0.foo0")
 
 
 def test_track_from_multiple_files(tmp_dir):
@@ -333,46 +336,51 @@ def test_track_from_multiple_files(tmp_dir):
     c = Context.load_from(tree, path2)
     context.merge_update(c)
 
-    def key_tracked(path, key):
-        return key in context.tracked[relpath(path)]
+    def key_tracked(d, path, key):
+        return key in d[relpath(path)]
 
-    with context.track():
+    with context.track() as tracked:
         context.select("Train")
-        assert not (key_tracked(path1, "Train") or key_tracked(path2, "Train"))
+        assert not (
+            key_tracked(tracked, path1, "Train")
+            or key_tracked(tracked, path2, "Train")
+        )
 
         context.select("Train.us")
         assert not (
-            key_tracked(path1, "Train.us") or key_tracked(path2, "Train.us")
+            key_tracked(tracked, path1, "Train.us")
+            or key_tracked(tracked, path2, "Train.us")
         )
 
         context.select("Train.us.lr")
-        assert key_tracked(path1, "Train.us.lr") and not key_tracked(
-            path2, "Train.us.lr"
+        assert key_tracked(tracked, path1, "Train.us.lr") and not key_tracked(
+            tracked, path2, "Train.us.lr"
         )
         context.select("Train.us.layers")
-        assert not key_tracked(path1, "Train.us.layers") and key_tracked(
-            path2, "Train.us.layers"
-        )
+        assert not key_tracked(
+            tracked, path1, "Train.us.layers"
+        ) and key_tracked(tracked, path2, "Train.us.layers")
 
     context = Context.clone(context)
-    assert not context.tracked
+    assert not context._tracked_data
 
     # let's see with an alias
     context["us"] = context["Train"]["us"]
-    with context.track():
+    with context.track() as tracked:
         context.select("us")
         assert not (
-            key_tracked(path1, "Train.us") or key_tracked(path2, "Train.us")
+            key_tracked(tracked, path1, "Train.us")
+            or key_tracked(tracked, path2, "Train.us")
         )
 
         context.select("us.lr")
-        assert key_tracked(path1, "Train.us.lr") and not key_tracked(
-            path2, "Train.us.lr"
+        assert key_tracked(tracked, path1, "Train.us.lr") and not key_tracked(
+            tracked, path2, "Train.us.lr"
         )
         context.select("Train.us.layers")
-        assert not key_tracked(path1, "Train.us.layers") and key_tracked(
-            path2, "Train.us.layers"
-        )
+        assert not key_tracked(
+            tracked, path1, "Train.us.layers"
+        ) and key_tracked(tracked, path2, "Train.us.layers")
 
 
 def test_node_value():
@@ -425,4 +433,4 @@ def test_resolve_resolves_boolean_value():
 def test_merge_from_raises_if_file_not_exist(tmp_dir, dvc):
     context = Context(foo="bar")
     with pytest.raises(ParamsFileNotFound):
-        context.merge_from(dvc.tree, tmp_dir / DEFAULT_PARAMS_FILE)
+        context.merge_from(dvc.tree, DEFAULT_PARAMS_FILE, tmp_dir)

tests/unit/test_interpolate.py

@@ -2,8 +2,7 @@ from math import inf, pi
 
 import pytest
 
-from dvc.parsing import Context
-from tests.func.test_stage_resolver import recurse_not_a_node
+from dvc.parsing.context import Context, recurse_not_a_node
 
 
 @pytest.mark.parametrize(
@@ -91,7 +90,7 @@ def test_resolve_primitives_list_access():
 
 
 def test_resolve_collection():
-    from .test_stage_resolver import (
+    from tests.func.parsing import (
         CONTEXT_DATA,
         RESOLVED_DVC_YAML_DATA,
         TEMPLATED_DVC_YAML_DATA,

tests/unit/test_stage_resolver.py

@@ -1,188 +0,0 @@
-from math import pi
-
-import pytest
-
-from dvc.parsing import DataResolver, ResolveError
-from dvc.parsing.context import Context, Value
-
-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(tmp_dir, dvc):
-    resolver = DataResolver(dvc, tmp_dir, TEMPLATED_DVC_YAML_DATA)
-    resolver.global_ctx = Context(CONTEXT_DATA)
-    assert resolver.resolve() == RESOLVED_DVC_YAML_DATA
-
-
-def test_set():
-    context = Context(CONTEXT_DATA)
-    to_set = {
-        "foo": "foo",
-        "bar": "bar",
-        "pi": pi,
-        "true": True,
-        "false": False,
-        "none": "None",
-        "int": 1,
-        "lst2": [1, 2, 3],
-        "dct2": {"foo": "bar", "foobar": "foobar"},
-    }
-    DataResolver.set_context_from(context, to_set)
-
-    for key, value in to_set.items():
-        # FIXME: using for convenience, figure out better way to do it
-        assert context[key] == context._convert(key, value)
-
-
-@pytest.mark.parametrize(
-    "coll",
-    [
-        ["foo", "bar", ["foo", "bar"]],
-        ["foo", "bar", {"foo": "foo", "bar": "bar"}],
-        {"foo": "foo", "bar": ["foo", "bar"]},
-        {"foo": "foo", "bar": {"foo": "foo", "bar": "bar"}},
-    ],
-)
-def test_set_nested_coll(coll):
-    context = Context(CONTEXT_DATA)
-    with pytest.raises(ResolveError) as exc_info:
-        DataResolver.set_context_from(context, {"thresh": 10, "item": coll})
-
-    assert (
-        str(exc_info.value) == "Failed to set 'item': Cannot set 'item', "
-        "has nested dict/list"
-    )
-
-
-def test_set_already_exists():
-    context = Context({"item": "foo"})
-    with pytest.raises(ResolveError) as exc_info:
-        DataResolver.set_context_from(context, {"item": "bar"})
-
-    assert (
-        str(exc_info.value) == "Failed to set 'item': Cannot set 'item', "
-        "key already exists"
-    )
-    assert context["item"] == Value("foo")
-
-
-@pytest.mark.parametrize(
-    "coll", [["foo", "${bar}"], {"foo": "${foo}", "bar": "bar"}],
-)
-def test_set_collection_interpolation(coll):
-    context = Context(CONTEXT_DATA)
-    with pytest.raises(ResolveError) as exc_info:
-        DataResolver.set_context_from(context, {"thresh": 10, "item": coll})
-
-    assert (
-        str(exc_info.value) == "Failed to set 'item': Cannot set 'item', "
-        f"having interpolation inside '{type(coll).__name__}' "
-        "is not supported."
-    )
-
-
-def test_set_interpolated_string():
-    context = Context(CONTEXT_DATA)
-    DataResolver.set_context_from(
-        context,
-        {
-            "foo": "${dict.foo}",
-            "bar": "${dict.bar}",
-            "param1": "${list[0]}",
-            "param2": "${list[1]}",
-            "frozen": "${freeze}",
-            "dict2": "${dict}",
-            "list2": "${list}",
-        },
-    )
-
-    assert context["foo"] == Value("foo")
-    assert context["bar"] == Value("bar")
-    assert context["param1"] == Value("param1")
-    assert context["param2"] == Value("param2")
-    assert context["frozen"] == context["freeze"] == Value(True)
-    assert context["dict2"] == context["dict"] == CONTEXT_DATA["dict"]
-    assert context["list2"] == context["list"] == CONTEXT_DATA["list"]
-
-
-def test_set_ladder():
-    context = Context(CONTEXT_DATA)
-    DataResolver.set_context_from(
-        context,
-        {
-            "item": 5,
-            "foo": "${dict.foo}",
-            "bar": "${dict.bar}",
-            "bar2": "${bar}",
-            "dict2": "${dict}",
-            "list2": "${list}",
-            "dict3": "${dict2}",
-            "list3": "${list2}",
-        },
-    )
-
-    assert context["item"] == Value(5)
-    assert context["foo"] == context["dict"]["foo"] == Value("foo")
-    assert (
-        context["bar"]
-        == context["bar2"]
-        == context["dict"]["bar"]
-        == Value("bar")
-    )
-    assert (
-        context["dict"]
-        == context["dict2"]
-        == context["dict3"]
-        == CONTEXT_DATA["dict"]
-    )
-    assert (
-        context["list"]
-        == context["list2"]
-        == context["list3"]
-        == CONTEXT_DATA["list"]
-    )
-
-
-@pytest.mark.parametrize(
-    "value",
-    ["param ${dict.foo}", "${dict.bar}${dict.foo}", "${dict.foo}-${dict.bar}"],
-)
-def test_set_multiple_interpolations(value):
-    context = Context(CONTEXT_DATA)
-    with pytest.raises(ResolveError,) as exc_info:
-        DataResolver.set_context_from(context, {"thresh": 10, "item": value})
-
-    assert str(exc_info.value) == (
-        "Failed to set 'item': Cannot set 'item', "
-        "joining string with interpolated string is not supported"
-    )