parametrization: Move resolver to StageLoader from Dvcfile (#4916)

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:

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}
 

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):

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

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