From 1c9588e8d503b4947dd5f01b64f1365bc1dbf225 Mon Sep 17 00:00:00 2001
From: nik123 <kodenkonikita@gmail.com>
Date: Tue, 14 Jul 2020 00:20:39 +0700
Subject: [PATCH] remove: remove dvc.yaml and dvc.lock if they are empty
 (#4196)

* remove: remove dvc.yaml, dvc.lock if they are empty

* tests: remove: simplified dvc.yaml in test_remove_stage_removes_dvcfiles_if_no_stages_left

* tests: remove: minor fixes

* tests: remove: simplified test_remove_stage_removes_dvcfiles_if_no_stages_left

* remove: independent  removal of dvc.lock and dvc.yaml
---
 dvc/dvcfile.py             | 11 +++++++++--
 tests/func/test_dvcfile.py | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py
index 8269edaa8a..5b0f2f38c6 100644
--- a/dvc/dvcfile.py
+++ b/dvc/dvcfile.py
@@ -233,7 +233,11 @@ def remove_stage(self, stage):
 
         logger.debug("Removing '%s' from '%s'", stage.name, self.path)
         del d["stages"][stage.name]
-        dump_yaml(self.path, d)
+
+        if d["stages"]:
+            dump_yaml(self.path, d)
+        else:
+            super().remove()
 
 
 class Lockfile(FileMixin):
@@ -286,7 +290,10 @@ def remove_stage(self, stage):
         logger.debug("Removing '%s' from '%s'", stage.name, self.path)
         del d[stage.name]
 
-        dump_yaml(self.path, d)
+        if d:
+            dump_yaml(self.path, d)
+        else:
+            self.remove()
 
 
 class Dvcfile:
diff --git a/tests/func/test_dvcfile.py b/tests/func/test_dvcfile.py
index 8c7978d648..5ba42a093c 100644
--- a/tests/func/test_dvcfile.py
+++ b/tests/func/test_dvcfile.py
@@ -285,6 +285,23 @@ def test_remove_stage_preserves_comment(tmp_dir, dvc, run_copy):
     )
 
 
+def test_remove_stage_removes_dvcfiles_if_no_stages_left(
+    tmp_dir, dvc, run_copy
+):
+    tmp_dir.gen("foo", "foo")
+    run_copy("foo", "bar", name="run_copy")
+
+    dvc_file = Dvcfile(dvc, PIPELINE_FILE)
+
+    assert dvc_file.exists()
+    assert (tmp_dir / PIPELINE_LOCK).exists()
+    assert (tmp_dir / "foo").exists()
+
+    dvc_file.remove_stage(dvc_file.stages["run_copy"])
+    assert not dvc_file.exists()
+    assert not (tmp_dir / PIPELINE_LOCK).exists()
+
+
 def test_dvcfile_dump_preserves_meta(tmp_dir, dvc, run_copy):
     tmp_dir.gen("foo", "foo")
     stage = run_copy("foo", "bar", name="run_copy")