From bc4c9fe9eabcd322996bf655a15582673d370b8e Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Tue, 24 Nov 2020 17:40:47 +0545
Subject: [PATCH] parametrization: prevent dumping parametrized stages (#4950)

* parametrization: prevent dumping parametrized stages

This is a bit hackish, but does the job done.
It feels like we need some unified way to be able to load
and dump, but not sure how. Till then, this dataclass
will be helpful in communicating between those.

* s/Meta/RawData

* Update dvc/stage/loader.py
---
 dvc/dvcfile.py             | 10 ++++++++++
 dvc/stage/__init__.py      |  9 +++++++++
 dvc/stage/loader.py        | 15 +++++++++++++--
 tests/func/test_dvcfile.py | 31 ++++++++++++++++++++++++++++++-
 4 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py
index 449799347..804f739c7 100644
--- a/dvc/dvcfile.py
+++ b/dvc/dvcfile.py
@@ -36,6 +36,10 @@ class LockfileCorruptedError(DvcException):
     pass
 
 
+class ParametrizedDumpError(DvcException):
+    pass
+
+
 def is_valid_filename(path):
     return path.endswith(DVC_FILE_SUFFIX) or os.path.basename(path) in [
         DVC_FILE,
@@ -197,7 +201,13 @@ class PipelineFile(FileMixin):
     def _dump_lockfile(self, stage):
         self._lockfile.dump(stage)
 
+    @staticmethod
+    def _check_if_parametrized(stage):
+        if stage.raw_data.parametrized:
+            raise ParametrizedDumpError(f"cannot dump a parametrized {stage}")
+
     def _dump_pipeline_file(self, stage):
+        self._check_if_parametrized(stage)
         stage_data = serialize.to_pipeline_file(stage)
 
         with modify_yaml(self.path, tree=self.repo.tree) as data:
diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py
index 839ba69f2..66dca0158 100644
--- a/dvc/stage/__init__.py
+++ b/dvc/stage/__init__.py
@@ -2,6 +2,8 @@ import logging
 import os
 import string
 from collections import defaultdict
+from dataclasses import dataclass
+from typing import Optional
 
 from funcy import cached_property, project
 
@@ -60,6 +62,12 @@ def loads_from(cls, repo, path, wdir, data):
     return cls(**kw)
 
 
+@dataclass
+class RawData:
+    parametrized: bool = False
+    generated_from: Optional[str] = None
+
+
 def create_stage(cls, repo, path, external=False, **kwargs):
     from dvc.dvcfile import check_dvc_filename
 
@@ -131,6 +139,7 @@ class Stage(params.StageParams):
         self._stage_text = stage_text
         self._dvcfile = dvcfile
         self.desc = desc
+        self.raw_data = RawData()
 
     @property
     def path(self) -> str:
diff --git a/dvc/stage/loader.py b/dvc/stage/loader.py
index acbc94a48..9b322c2ae 100644
--- a/dvc/stage/loader.py
+++ b/dvc/stage/loader.py
@@ -7,7 +7,7 @@ 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.parsing import JOIN, DataResolver
 from dvc.path_info import PathInfo
 
 from . import PipelineStage, Stage, loads_from
@@ -22,6 +22,7 @@ class StageLoader(Mapping):
     def __init__(self, dvcfile, data, lockfile_data=None):
         self.dvcfile = dvcfile
         self.data = data or {}
+        self.stages_data = self.data.get("stages", {})
         self.repo = self.dvcfile.repo
         self._enable_parametrization = self.repo.config["feature"][
             "parametrization"
@@ -106,13 +107,23 @@ class StageLoader(Mapping):
                 "No lock entry found for '%s:%s'", self.dvcfile.relpath, name,
             )
 
-        return self.load_stage(
+        stage = self.load_stage(
             self.dvcfile,
             name,
             self.resolved_data[name],
             self.lockfile_data.get(name, {}),
         )
 
+        group, *keys = name.rsplit(JOIN, maxsplit=1)
+        if group and keys and name not in self.stages_data:
+            stage.raw_data.generated_from = group
+
+        stage.raw_data.parametrized = (
+            self.stages_data.get(name, {}) != self.resolved_data[name]
+        )
+
+        return stage
+
     def __iter__(self):
         return iter(self.resolved_data)
 
diff --git a/tests/func/test_dvcfile.py b/tests/func/test_dvcfile.py
index 82b708b2b..f11f960bb 100644
--- a/tests/func/test_dvcfile.py
+++ b/tests/func/test_dvcfile.py
@@ -3,7 +3,13 @@ import textwrap
 
 import pytest
 
-from dvc.dvcfile import PIPELINE_FILE, PIPELINE_LOCK, Dvcfile, SingleStageFile
+from dvc.dvcfile import (
+    PIPELINE_FILE,
+    PIPELINE_LOCK,
+    Dvcfile,
+    ParametrizedDumpError,
+    SingleStageFile,
+)
 from dvc.stage.exceptions import (
     StageFileDoesNotExistError,
     StageFileFormatError,
@@ -355,3 +361,26 @@ def test_dvcfile_dump_preserves_comments(tmp_dir, dvc):
 
     dvcfile.dump(stage)
     assert dvcfile._load()[1] == (text + ":\n\tcache: false\n".expandtabs())
+
+
+@pytest.mark.parametrize(
+    "data, name",
+    [
+        ({"build-us": {"cmd": "echo ${foo}"}}, "build-us"),
+        (
+            {"build": {"foreach": ["us", "gb"], "do": {"cmd": "echo ${foo}"}}},
+            "build@us",
+        ),
+    ],
+)
+def test_dvcfile_try_dumping_parametrized_stage(tmp_dir, dvc, data, name):
+    dump_yaml("dvc.yaml", {"stages": data, "vars": [{"foo": "foobar"}]})
+    dvc.config["feature"]["parametrization"] = True
+
+    stage = dvc.get_stage(name=name)
+    dvcfile = stage.dvcfile
+
+    with pytest.raises(ParametrizedDumpError) as exc:
+        dvcfile.dump(stage)
+
+    assert str(exc.value) == f"cannot dump a parametrized stage: '{name}'"