diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py index 1a8e42fe4..fecc209b1 100644 --- a/dvc/parsing/__init__.py +++ b/dvc/parsing/__init__.py @@ -136,7 +136,7 @@ class DataResolver: suffix = str(key if key is not DEFAULT_SENTINEL else value) return self._resolve_stage(c, f"{name}-{suffix}", in_data) - iterable = context.resolve(foreach_data) + iterable = context.resolve(foreach_data, unwrap=False) assert isinstance(iterable, (Sequence, Mapping)) and not isinstance( iterable, str diff --git a/dvc/parsing/context.py b/dvc/parsing/context.py index ae35ca80d..4cc2848b3 100644 --- a/dvc/parsing/context.py +++ b/dvc/parsing/context.py @@ -1,4 +1,5 @@ import os +from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Mapping, MutableMapping, MutableSequence, Sequence from contextlib import contextmanager @@ -6,7 +7,7 @@ from copy import deepcopy from dataclasses import dataclass, field, replace from typing import Any, List, Optional, Union -from funcy import identity +from funcy import identity, rpartial from dvc.parsing.interpolate import ( get_expression, @@ -30,6 +31,7 @@ def _merge(into, update, overwrite): f"Cannot overwrite as key {key} already exists in {into}" ) into[key] = val + assert isinstance(into[key], Node) @dataclass @@ -55,24 +57,41 @@ def _default_meta(): return Meta(source=None) +class Node: + def get_sources(self): + raise NotImplementedError + + @property + @abstractmethod + def value(self): + pass + + @dataclass -class Value: - value: Any +class Value(Node): + _value: Any meta: Meta = field( compare=False, default_factory=_default_meta, repr=False ) def __repr__(self): - return repr(self.value) + return repr(self._value) def __str__(self) -> str: - return str(self.value) + return str(self._value) def get_sources(self): return {self.meta.source: self.meta.path()} + @property + def value(self): + return self._value -class Container: + +PRIMITIVES = (int, float, str, bytes, bool) + + +class Container(Node, ABC): meta: Meta data: Union[list, dict] _key_transform = staticmethod(identity) @@ -82,9 +101,9 @@ class Container: def _convert(self, key, value): meta = Meta.update_path(self.meta, key) - if value is None or isinstance(value, (int, float, str, bytes, bool)): + if value is None or isinstance(value, PRIMITIVES): return Value(value, meta=meta) - elif isinstance(value, (CtxList, CtxDict, Value)): + elif isinstance(value, Node): return value elif isinstance(value, (list, dict)): container = CtxDict if isinstance(value, dict) else CtxList @@ -150,6 +169,10 @@ class CtxList(Container, MutableSequence): def get_sources(self): return {self.meta.source: self.meta.path()} + @property + def value(self): + return [node.value for node in self] + class CtxDict(Container, MutableMapping): def __init__(self, mapping: Mapping = None, meta: Meta = None, **kwargs): @@ -169,7 +192,11 @@ class CtxDict(Container, MutableMapping): def merge_update(self, *args, overwrite=False): for d in args: - _merge(self.data, d, overwrite=overwrite) + _merge(self, d, overwrite=overwrite) + + @property + def value(self): + return {key: node.value for key, node in self.items()} class Context(CtxDict): @@ -205,10 +232,21 @@ class Context(CtxDict): def select( self, key: str, unwrap=False ): # pylint: disable=arguments-differ - # NOTE: `unwrap` default value is different from `resolve_str` + """Select the item using key, similar to `__getitem__` + but can track the usage of the data on interpolation + as well and can get from nested data structure by using + "." separated key (eg: "key1.key2.key3") + + Args: + key: key to select value from + unwrap: Convert CtxList/CtxDict/Value items to it's original data + Defaults to False. Note that the default is different from + `resolve`. + """ node = super().select(key) self._track_data(node) - return node.value if isinstance(node, Value) and unwrap else node + assert isinstance(node, Node) + return node.value if unwrap else node @classmethod def load_from(cls, tree, file: str) -> "Context": @@ -243,20 +281,27 @@ class Context(CtxDict): self._check_interpolation_collection(key, value) self[key] = value - def resolve(self, src): + def resolve(self, src, unwrap=True): """Recursively resolves interpolation and returns resolved data. + Args: + src: Data (str/list/dict etc.) to resolve + unwrap: Unwrap CtxDict/CtxList/Value to it's original data if + inside `src`. Defaults to True. + >>> c = Context({"three": 3}) >>> c.resolve({"lst": [1, 2, "${three}"]}) {'lst': [1, 2, 3]} """ Seq = (list, tuple, set) + + resolve = rpartial(self.resolve, unwrap) if isinstance(src, Mapping): - return {key: self.resolve(value) for key, value in src.items()} + return {key: resolve(value) for key, value in src.items()} elif isinstance(src, Seq): - return type(src)(map(self.resolve, src)) + return type(src)(map(resolve, src)) elif isinstance(src, str): - return self.resolve_str(src) + return self.resolve_str(src, unwrap=unwrap) return src def resolve_str(self, src: str, unwrap=True): diff --git a/tests/func/test_stage_resolver.py b/tests/func/test_stage_resolver.py index c91d8d002..28e4c6e12 100644 --- a/tests/func/test_stage_resolver.py +++ b/tests/func/test_stage_resolver.py @@ -6,6 +6,7 @@ import pytest from dvc.dependency import _merge_params from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver +from dvc.parsing.context import Node from dvc.path_info import PathInfo from dvc.utils.serialize import dump_json, dump_yaml @@ -44,9 +45,19 @@ RESOLVED_DVC_YAML_DATA = { } +def recurse_not_a_node(d): + assert not isinstance(d, Node) + if isinstance(d, (list, dict)): + iterable = d if isinstance(d, list) else d.values() + for item in iterable: + assert recurse_not_a_node(item) + return True + + def assert_stage_equal(d1, d2): """Keeps the params section in order, and then checks for equality.""" for d in [d1, d2]: + assert recurse_not_a_node(d) for _, stage_d in d.get("stages", {}).items(): params = _merge_params(stage_d.get("params", [])) for k in params: @@ -248,12 +259,15 @@ def test_simple_foreach_loop(tmp_dir, dvc): } resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d) - assert resolver.resolve() == { - "stages": { - f"build-{item}": {"cmd": f"python script.py {item}"} - for item in iterable - } - } + assert_stage_equal( + resolver.resolve(), + { + "stages": { + f"build-{item}": {"cmd": f"python script.py {item}"} + for item in iterable + } + }, + ) def test_foreach_loop_dict(tmp_dir, dvc): @@ -268,12 +282,15 @@ def test_foreach_loop_dict(tmp_dir, dvc): } resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d) - assert resolver.resolve() == { - "stages": { - f"build-{key}": {"cmd": f"python script.py {item['thresh']}"} - for key, item in iterable["models"].items() - } - } + assert_stage_equal( + resolver.resolve(), + { + "stages": { + f"build-{key}": {"cmd": f"python script.py {item['thresh']}"} + for key, item in iterable["models"].items() + } + }, + ) def test_foreach_loop_templatized(tmp_dir, dvc): @@ -319,14 +336,17 @@ def test_set(tmp_dir, dvc, value): } } resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d) - assert resolver.resolve() == { - "stages": { - "build": { - "cmd": f"python script.py --thresh {value}", - "always_changed": value, + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "build": { + "cmd": f"python script.py --thresh {value}", + "always_changed": value, + } } - } - } + }, + ) @pytest.mark.parametrize( @@ -343,11 +363,14 @@ def test_coll(tmp_dir, dvc, coll): } } resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d) - assert resolver.resolve() == { - "stages": { - "build": {"cmd": "python script.py --thresh 10", "outs": coll} - } - } + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "build": {"cmd": "python script.py --thresh 10", "outs": coll} + } + }, + ) def test_set_with_foreach(tmp_dir, dvc): @@ -362,12 +385,15 @@ def test_set_with_foreach(tmp_dir, dvc): } } resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d) - assert resolver.resolve() == { - "stages": { - f"build-{item}": {"cmd": f"command --value {item}"} - for item in items - } - } + assert_stage_equal( + resolver.resolve(), + { + "stages": { + f"build-{item}": {"cmd": f"command --value {item}"} + for item in items + } + }, + ) def test_set_with_foreach_and_on_stage_definition(tmp_dir, dvc): @@ -388,15 +414,18 @@ def test_set_with_foreach_and_on_stage_definition(tmp_dir, dvc): }, } resolver = DataResolver(dvc, PathInfo(str(tmp_dir)), d) - assert resolver.resolve() == { - "stages": { - "build-us": { - "cmd": "command --value 10", - "params": [{"params.json": ["models.us.thresh"]}], - }, - "build-gb": { - "cmd": "command --value 15", - "params": [{"params.json": ["models.gb.thresh"]}], - }, - } - } + assert_stage_equal( + resolver.resolve(), + { + "stages": { + "build-us": { + "cmd": "command --value 10", + "params": [{"params.json": ["models.us.thresh"]}], + }, + "build-gb": { + "cmd": "command --value 15", + "params": [{"params.json": ["models.gb.thresh"]}], + }, + } + }, + ) diff --git a/tests/unit/test_context.py b/tests/unit/test_context.py index 6d8da1d2c..2d4e20de4 100644 --- a/tests/unit/test_context.py +++ b/tests/unit/test_context.py @@ -6,6 +6,7 @@ import pytest from dvc.parsing.context import Context, CtxDict, CtxList, Value from dvc.tree.local import LocalTree from dvc.utils.serialize import dump_yaml +from tests.func.test_stage_resolver import recurse_not_a_node def test_context(): @@ -145,6 +146,22 @@ def test_select(): assert context.select(f"lst.{i}.bar{i}") == Value(f"bar{i}") +def test_select_unwrap(): + context = Context({"dct": {"foo": "bar"}}, lst=[1, 2, 3], foo="foo") + + assert context.select("dct.foo", unwrap=True) == "bar" + assert context.select("lst.0", unwrap=True) == 1 + assert context.select("foo", unwrap=True) == "foo" + + node = context.select("dct", unwrap=True) + assert isinstance(node, dict) and recurse_not_a_node(node) + assert node == {"foo": "bar"} + + node = context.select("lst", unwrap=True) + assert isinstance(node, list) and recurse_not_a_node(node) + assert node == [1, 2, 3] + + def test_merge_dict(): d1 = {"Train": {"us": {"lr": 10}}} d2 = {"Train": {"us": {"layers": 100}}} @@ -160,7 +177,11 @@ def test_merge_dict(): c1.merge_update({"Train": {"us": {"lr": 15}}}) c1.merge_update({"Train": {"us": {"lr": 15}}}, overwrite=True) - assert c1.select("Train.us") == CtxDict(lr=15, layers=100) + node = c1.select("Train.us") + assert node == {"lr": 15, "layers": 100} + assert isinstance(node, CtxDict) + assert node["lr"] == Value(15) + assert node["layers"] == Value(100) def test_merge_list(): @@ -171,7 +192,10 @@ def test_merge_list(): # lists are never merged c1.merge_update({"lst": [10, 11, 12]}, overwrite=True) - assert c1.select("lst") == [10, 11, 12] + node = c1.select("lst") + assert node == [10, 11, 12] + assert isinstance(node, CtxList) + assert node[0] == Value(10) def test_overwrite_with_setitem(): @@ -319,3 +343,30 @@ def test_track_from_multiple_files(tmp_dir): assert not key_tracked(path1, "Train.us.layers") and key_tracked( path2, "Train.us.layers" ) + + +def test_node_value(): + d = {"dct": {"foo": "bar"}, "lst": [1, 2, 3], "foo": "foo"} + context = Context(d) + assert isinstance(context, (Context, CtxDict)) + assert isinstance(context["dct"], CtxDict) + assert isinstance(context["lst"], CtxList) + assert isinstance(context["foo"], Value) + assert isinstance(context["dct"]["foo"], Value) + assert isinstance(context["lst"][0], Value) + + assert context.value == d + assert recurse_not_a_node(context.value) + assert isinstance(context.value["dct"], dict) + assert isinstance(context.value["lst"], list) + assert isinstance(context.value["foo"], str) + assert isinstance(context.value["dct"]["foo"], str) + assert isinstance(context.value["lst"][0], int) + + assert isinstance(context["dct"].value, dict) + assert context["dct"]["foo"].value == "bar" + + assert isinstance(context["lst"].value, list) + assert context["lst"][1].value == 2 + + assert context["foo"].value == "foo" diff --git a/tests/unit/test_interpolate.py b/tests/unit/test_interpolate.py index 007dba495..7a8f1ebdf 100644 --- a/tests/unit/test_interpolate.py +++ b/tests/unit/test_interpolate.py @@ -3,6 +3,7 @@ from math import inf, pi import pytest from dvc.parsing import Context +from tests.func.test_stage_resolver import recurse_not_a_node @pytest.mark.parametrize( @@ -97,4 +98,6 @@ def test_resolve_collection(): ) context = Context(CONTEXT_DATA) - assert context.resolve(TEMPLATED_DVC_YAML_DATA) == RESOLVED_DVC_YAML_DATA + resolved = context.resolve(TEMPLATED_DVC_YAML_DATA) + assert resolved == RESOLVED_DVC_YAML_DATA + assert recurse_not_a_node(resolved)