diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py index 1e5168dde..b173de438 100644 --- a/dvc/parsing/__init__.py +++ b/dvc/parsing/__init__.py @@ -1,5 +1,6 @@ import logging import os +from collections import defaultdict from copy import deepcopy from itertools import starmap from typing import TYPE_CHECKING @@ -30,6 +31,7 @@ class DataResolver: def __init__(self, repo: "Repo", wdir: PathInfo, d: dict): to_import: PathInfo = wdir / d.get(USE_KWD, DEFAULT_PARAMS_FILE) vars_ = d.get(VARS_KWD, {}) + vars_ctx = Context(vars_) if os.path.exists(to_import): self.global_ctx_source = to_import self.global_ctx = Context.load_from(repo.tree, str(to_import)) @@ -41,7 +43,7 @@ class DataResolver: to_import, ) - self.global_ctx.merge_update(vars_) + self.global_ctx.merge_update(vars_ctx) self.data: dict = d self.wdir = wdir self.repo = repo @@ -92,12 +94,21 @@ class DataResolver: with context.track(): stage_d = resolve(definition, context) - params = stage_d.get(PARAMS_KWD, []) + context.tracked + params = stage_d.get(PARAMS_KWD, []) + self._resolve_params( + context, wdir + ) if params: stage_d[PARAMS_KWD] = params return {name: stage_d} + def _resolve_params(self, context: Context, wdir): + tracked = defaultdict(set) + for src, keys in context.tracked.items(): + tracked[str(PathInfo(src).relative_to(wdir))].update(keys) + + return [{file: list(keys)} for file, keys in tracked.items()] + def _resolve_wdir(self, context: Context, wdir: str = None) -> PathInfo: if not wdir: return self.wdir diff --git a/dvc/parsing/context.py b/dvc/parsing/context.py index 8211f5d88..5817c46b2 100644 --- a/dvc/parsing/context.py +++ b/dvc/parsing/context.py @@ -185,9 +185,7 @@ class Context(CtxDict): @property def tracked(self): - return [ - {file: list(keys)} for file, keys in self._tracked_data.items() - ] + return self._tracked_data def select(self, key: str): node = super().select(key) diff --git a/tests/func/test_stage_resolver.py b/tests/func/test_stage_resolver.py new file mode 100644 index 000000000..2c36aac9d --- /dev/null +++ b/tests/func/test_stage_resolver.py @@ -0,0 +1,235 @@ +import os +from copy import deepcopy + +import pytest + +from dvc.dependency import _merge_params +from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver +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"], + "params": [{"params.yaml": ["dict.foo", "dict.bar", "freeze"]}], + "frozen": True, + }, + "stage2": { + "cmd": "echo foo bar", + "params": [{"params.yaml": ["dict.foo", "dict.bar"]}], + }, + } +} + + +def assert_stage_equal(d1, d2): + """Keeps the params section in order, and then checks for equality.""" + for d in [d1, d2]: + 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)) + + +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) + + # `vars` section is not auto-tracked + del resolved_data["stages"]["stage1"]["params"] + del resolved_data["stages"]["stage2"]["params"] + assert_stage_equal(resolver.resolve(), resolved_data) + + +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(ValueError): + resolver.resolve() + + +def test_use(tmp_dir, dvc): + """ + Test that different file can be loaded using `use` + instead of default params.yaml. + """ + dump_yaml(tmp_dir / "params2.yaml", CONTEXT_DATA) + d = deepcopy(TEMPLATED_DVC_YAML_DATA) + d["use"] = "params2.yaml" + resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d) + + resolved_data = deepcopy(RESOLVED_DVC_YAML_DATA) + for _, stage_d in resolved_data["stages"].items(): + params = stage_d["params"][0] + params["params2.yaml"] = params.pop(DEFAULT_PARAMS_FILE) + + assert_stage_equal(resolver.resolve(), resolved_data) + + +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 = { + "use": DEFAULT_PARAMS_FILE, + "vars": {"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", "params": ["dict.bar"]} + } + }, + ) + + +def test_with_params_section(tmp_dir, dvc): + """Test that params section is also loaded for interpolation""" + d = { + "use": "params.yaml", + "vars": {"dict": {"foo": "foo"}}, + "stages": { + "stage1": { + "cmd": "echo ${dict.foo} ${dict.bar} ${dict.foobar}", + "params": [{"params.json": ["value1"]}], + }, + }, + } + 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": [ + "dict.bar", + {"params.json": ["dict.foobar", "value1"]}, + ], + } + } + }, + ) + + +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", + }, + }, + } + + 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) + + root_params_file = os.path.relpath(tmp_dir / "params.yaml", data_dir) + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "stage1": { + "cmd": "echo foo bar", + "wdir": "data", + "params": [ + "dict.foo", + "value1", + {root_params_file: ["dict.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}", + }, + }, + } + 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) + + root_params_file = os.path.relpath(tmp_dir / "params.yaml", data_dir) + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "stage1": { + "cmd": "echo foo bar", + "wdir": "data", + "params": [ + "dict.foo", + "value1", + {root_params_file: ["dict.bar", "dict.ws"]}, + ], + } + } + }, + ) diff --git a/tests/unit/test_context.py b/tests/unit/test_context.py index cd8e83863..6d8da1d2c 100644 --- a/tests/unit/test_context.py +++ b/tests/unit/test_context.py @@ -1,9 +1,7 @@ -from collections import defaultdict from dataclasses import asdict from math import pi import pytest -from funcy.seqs import first from dvc.parsing.context import Context, CtxDict, CtxList, Value from dvc.tree.local import LocalTree @@ -244,7 +242,7 @@ def test_track(tmp_dir): def key_tracked(key): assert len(context.tracked) == 1 - return key in context.tracked[0][str(path)] + return key in context.tracked[str(path)] with context.track(): context.select("lst") @@ -282,11 +280,7 @@ def test_track_from_multiple_files(tmp_dir): context.merge_update(c) def key_tracked(path, key): - tracked = defaultdict(list) - for item in context.tracked: - source = first(item) - tracked[source].extend(item[source]) - return key in tracked[str(path)] + return key in context.tracked[str(path)] with context.track(): context.select("Train")