From d5312e56a186b6d207a805ad2f8400352f66a0c7 Mon Sep 17 00:00:00 2001 From: Saugat Pachhai Date: Fri, 20 Nov 2020 13:38:57 +0545 Subject: [PATCH] parametrization: Move resolver to StageLoader from `Dvcfile` (#4916) --- dvc/dvcfile.py | 13 +------- dvc/parsing/__init__.py | 2 +- dvc/stage/loader.py | 30 ++++++++++++++----- tests/func/test_feature_flags.py | 15 ++++------ tests/unit/stage/test_loader_pipeline_file.py | 4 ++- 5 files changed, 33 insertions(+), 31 deletions(-) diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py index 2a46237d4..449799347 100644 --- a/dvc/dvcfile.py +++ b/dvc/dvcfile.py @@ -3,12 +3,9 @@ import contextlib import logging import os -from funcy import log_durations from voluptuous import MultipleInvalid from dvc.exceptions import DvcException -from dvc.parsing import DataResolver -from dvc.path_info import PathInfo from dvc.stage import serialize from dvc.stage.exceptions import ( StageFileBadNameError, @@ -233,16 +230,8 @@ class PipelineFile(FileMixin): @property def stages(self): data, _ = self._load() - - if self.repo.config["feature"]["parametrization"]: - with log_durations(logger.debug, "resolving values"): - resolver = DataResolver( - self.repo, PathInfo(self.path).parent, data - ) - data = resolver.resolve() - lockfile_data = self._lockfile.load() - return StageLoader(self, data.get("stages", {}), lockfile_data) + return StageLoader(self, data, lockfile_data) def remove(self, force=False): if not force: diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py index 914e011e3..1fad2095c 100644 --- a/dvc/parsing/__init__.py +++ b/dvc/parsing/__init__.py @@ -84,7 +84,7 @@ class DataResolver: def resolve(self): stages = self.data.get(STAGES_KWD, {}) - data = join(starmap(self._resolve_entry, stages.items())) + data = join(starmap(self._resolve_entry, stages.items())) or {} logger.trace("Resolved dvc.yaml:\n%s", data) return {STAGES_KWD: data} diff --git a/dvc/stage/loader.py b/dvc/stage/loader.py index 5f952aa93..acbc94a48 100644 --- a/dvc/stage/loader.py +++ b/dvc/stage/loader.py @@ -3,10 +3,12 @@ from collections.abc import Mapping from copy import deepcopy from itertools import chain -from funcy import get_in, lcat, project +from funcy import cached_property, get_in, lcat, log_durations, project from dvc import dependency, output from dvc.hash_info import HashInfo +from dvc.parsing import DataResolver +from dvc.path_info import PathInfo from . import PipelineStage, Stage, loads_from from .exceptions import StageNameUnspecified, StageNotFound @@ -17,11 +19,25 @@ logger = logging.getLogger(__name__) class StageLoader(Mapping): - def __init__(self, dvcfile, stages_data, lockfile_data=None): + def __init__(self, dvcfile, data, lockfile_data=None): self.dvcfile = dvcfile - self.stages_data = stages_data or {} + self.data = data or {} + self.repo = self.dvcfile.repo + self._enable_parametrization = self.repo.config["feature"][ + "parametrization" + ] self.lockfile_data = lockfile_data or {} + @cached_property + def resolved_data(self): + data = self.data + if self._enable_parametrization: + wdir = PathInfo(self.dvcfile.path).parent + with log_durations(logger.debug, "resolving values"): + resolver = DataResolver(self.repo, wdir, data) + data = 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.""" @@ -93,18 +109,18 @@ class StageLoader(Mapping): return self.load_stage( self.dvcfile, name, - self.stages_data[name], + self.resolved_data[name], self.lockfile_data.get(name, {}), ) def __iter__(self): - return iter(self.stages_data) + return iter(self.resolved_data) def __len__(self): - return len(self.stages_data) + return len(self.resolved_data) def __contains__(self, name): - return name in self.stages_data + return name in self.resolved_data class SingleStageLoader(Mapping): diff --git a/tests/func/test_feature_flags.py b/tests/func/test_feature_flags.py index 40c900413..cc59cbdff 100644 --- a/tests/func/test_feature_flags.py +++ b/tests/func/test_feature_flags.py @@ -1,4 +1,5 @@ from dvc.dvcfile import Dvcfile +from dvc.stage.loader import StageLoader from dvc.utils.serialize import dumps_yaml from tests.unit.test_stage_resolver import ( RESOLVED_DVC_YAML_DATA, @@ -10,7 +11,7 @@ def test_parametrization_is_not_enabled_by_default(tmp_dir, dvc, mocker): assert dvc.config["feature"]["parametrization"] is False (tmp_dir / "dvc.yaml").write_text(dumps_yaml(RESOLVED_DVC_YAML_DATA)) - mock = mocker.patch("dvc.dvcfile.DataResolver") + mock = mocker.patch("dvc.stage.loader.DataResolver") stages = list(Dvcfile(dvc, "dvc.yaml").stages) mock.assert_not_called() @@ -20,15 +21,9 @@ def test_parametrization_is_not_enabled_by_default(tmp_dir, dvc, mocker): def test_parametrization_flag_when_enabled(tmp_dir, dvc, mocker): dvc.config["feature"]["parametrization"] = True - mock = mocker.patch( - "dvc.dvcfile.DataResolver.resolve", return_value=RESOLVED_DVC_YAML_DATA - ) - dvcfile = Dvcfile(dvc, "dvc.yaml") - mocker.patch.object( - dvcfile, "_load", return_value=[TEMPLATED_DVC_YAML_DATA, None] - ) + loader = StageLoader(dvcfile, TEMPLATED_DVC_YAML_DATA, lockfile_data=None) + loader.resolved_data = RESOLVED_DVC_YAML_DATA["stages"] - stages = list(dvcfile.stages) - mock.assert_called_once() + stages = list(loader.values()) assert len(stages) == 2 diff --git a/tests/unit/stage/test_loader_pipeline_file.py b/tests/unit/stage/test_loader_pipeline_file.py index dad059bc8..a2c476a22 100644 --- a/tests/unit/stage/test_loader_pipeline_file.py +++ b/tests/unit/stage/test_loader_pipeline_file.py @@ -222,7 +222,9 @@ def test_load_stage_wdir_and_path_correctly(dvc, stage_data, lock_data): def test_load_stage_mapping(dvc, stage_data, lock_data): dvcfile = Dvcfile(dvc, PIPELINE_FILE) - loader = StageLoader(dvcfile, {"stage": stage_data}, {"stage": lock_data}) + loader = StageLoader( + dvcfile, {"stages": {"stage": stage_data}}, {"stage": lock_data} + ) assert len(loader) == 1 assert "stage" in loader assert "stage1" not in loader