From 24a14563cc35e34eb27671d96b45b72711dc2a34 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Fri, 12 Jun 2020 21:55:15 +0000
Subject: [PATCH] run: preserve order when -f/--force overwrite stage (#4035)

This commit preserves comments, meta and order inside a dvc.yaml
when `-f/--force` is used. But, it does not maintain comments/meta
for .dvc files.

Related: #4014
---
 dvc/repo/run.py                     | 11 +++----
 tests/func/test_run_multistage.py   | 46 +++++++++++++++++++++++++++++
 tests/func/test_run_single_stage.py | 30 +++++++++++++++++++
 3 files changed, 82 insertions(+), 5 deletions(-)

diff --git a/dvc/repo/run.py b/dvc/repo/run.py
index a5bb81f46..6cf8d8e72 100644
--- a/dvc/repo/run.py
+++ b/dvc/repo/run.py
@@ -1,4 +1,5 @@
 import os
+from contextlib import suppress
 
 from funcy import concat, first, lfilter
 
@@ -109,12 +110,12 @@ def run(self, fname=None, no_exec=False, single_stage=False, **kwargs):
         return None
 
     dvcfile = Dvcfile(self, stage.path)
-    if kwargs.get("force", True):
-        dvcfile.remove_stage(stage)
-    else:
-        _check_stage_exists(dvcfile, stage)
-
     try:
+        if kwargs.get("force", True):
+            with suppress(ValueError):
+                self.stages.remove(stage)
+        else:
+            _check_stage_exists(dvcfile, stage)
         self.check_modified_graph([stage])
     except OutputDuplicationError as exc:
         raise OutputDuplicationError(exc.output, set(exc.stages) - {stage})
diff --git a/tests/func/test_run_multistage.py b/tests/func/test_run_multistage.py
index cafb65f86..c300f1cee 100644
--- a/tests/func/test_run_multistage.py
+++ b/tests/func/test_run_multistage.py
@@ -1,4 +1,5 @@
 import os
+import textwrap
 
 import pytest
 import yaml
@@ -6,6 +7,7 @@ import yaml
 from dvc.exceptions import InvalidArgumentError
 from dvc.repo import Repo
 from dvc.stage.exceptions import DuplicateStageName, InvalidStageName
+from dvc.utils.yaml import parse_yaml_for_update
 
 
 def test_run_with_name(tmp_dir, dvc, run_copy):
@@ -315,3 +317,47 @@ def test_run_without_cmd(kwargs):
     with pytest.raises(InvalidArgumentError) as exc:
         Repo().run(**kwargs)
     assert "command is not specified" == str(exc.value)
+
+
+def test_run_overwrite_order(tmp_dir, dvc, run_copy):
+    from dvc.dvcfile import PIPELINE_FILE
+
+    tmp_dir.gen({"foo": "foo", "foo1": "foo1"})
+    run_copy("foo", "bar", name="copy-foo-bar")
+    run_copy("bar", "foobar", name="copy-bar-foobar")
+
+    run_copy("foo1", "bar1", name="copy-foo-bar", force=True)
+
+    data = parse_yaml_for_update(
+        (tmp_dir / PIPELINE_FILE).read_text(), PIPELINE_FILE
+    )
+    assert list(data["stages"].keys()) == ["copy-foo-bar", "copy-bar-foobar"]
+
+
+def test_run_overwrite_preserves_meta_and_comment(tmp_dir, dvc, run_copy):
+    from dvc.dvcfile import PIPELINE_FILE
+
+    tmp_dir.gen({"foo": "foo", "foo1": "foo1"})
+    text = textwrap.dedent(
+        """\
+        stages:
+          copy-foo-bar:
+            cmd: python copy.py {src} {dest}
+            deps:
+            - copy.py
+            - {src}
+            outs:
+            # comments are preserved
+            - {dest}
+            meta:
+              name: meta is preserved too
+    """
+    )
+    (tmp_dir / PIPELINE_FILE).write_text(text.format(src="foo", dest="bar"))
+    assert dvc.reproduce(PIPELINE_FILE)
+
+    assert run_copy("foo1", "bar1", name="copy-foo-bar", force=True)
+
+    assert (tmp_dir / PIPELINE_FILE).read_text() == text.format(
+        src="foo1", dest="bar1"
+    )
diff --git a/tests/func/test_run_single_stage.py b/tests/func/test_run_single_stage.py
index 48df72147..221fee116 100644
--- a/tests/func/test_run_single_stage.py
+++ b/tests/func/test_run_single_stage.py
@@ -1,6 +1,7 @@
 import filecmp
 import logging
 import os
+import textwrap
 import uuid
 from pathlib import Path
 
@@ -1017,3 +1018,32 @@ class TestRunDirMetrics:
                 metrics_no_cache=["dir"],
                 single_stage=True,
             )
+
+
+def test_run_force_doesnot_preserve_comments_and_meta(tmp_dir, dvc, run_copy):
+    """Depends on loading of stage on `run` where we don't check the file
+    for stage already exists, so we don't copy `stage_text` over due to which
+    `meta` and `comments` don't get preserved."""
+    tmp_dir.gen({"foo": "foo", "foo1": "foo1"})
+    text = textwrap.dedent(
+        """\
+      cmd: python copy.py foo bar
+      deps:
+      - path: copy.py
+      - path: foo
+      outs:
+      # comment not preserved
+      - path: bar
+      meta:
+        name: copy-foo-bar
+    """
+    )
+    (tmp_dir / "bar.dvc").write_text(text)
+    dvc.reproduce("bar.dvc")
+    assert "comment" in (tmp_dir / "bar.dvc").read_text()
+    assert "meta" in (tmp_dir / "bar.dvc").read_text()
+
+    run_copy("foo1", "bar1", single_stage=True, force=True, fname="bar.dvc")
+
+    assert "comment" not in (tmp_dir / "bar.dvc").read_text()
+    assert "meta" not in (tmp_dir / "bar.dvc").read_text()