From 665185b9e7491433b39ccc61398fa8eff02659f6 Mon Sep 17 00:00:00 2001
From: David de la Iglesia Castro <daviddelaiglesiacastro@gmail.com>
Date: Thu, 24 Nov 2022 10:38:18 +0100
Subject: [PATCH] exp: save: Handle staged changes.

Staged changes were causing a merge error.
Warn and unstage instead of erroring, to match exp run behavior.
---
 dvc/repo/experiments/save.py        | 10 +++++++++-
 tests/func/experiments/test_save.py | 12 +++++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/dvc/repo/experiments/save.py b/dvc/repo/experiments/save.py
index 2785ca07e..692a4b6c7 100644
--- a/dvc/repo/experiments/save.py
+++ b/dvc/repo/experiments/save.py
@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, List, Optional
 
 from funcy import first
 
@@ -23,6 +23,14 @@ def save(
     queue = repo.experiments.workspace_queue
     logger.debug("Saving workspace in %s", os.getcwd())
 
+    staged, _, _ = repo.scm.status(untracked_files="no")
+    if staged:
+        logger.warning(
+            "Your workspace contains staged Git changes which will be "
+            "unstaged before saving this experiment."
+        )
+        repo.scm.reset()
+
     entry = repo.experiments.new(queue=queue, name=name, force=force)
     executor = queue.init_executor(repo.experiments, entry)
     save_result = executor.save(executor.info, force=force)
diff --git a/tests/func/experiments/test_save.py b/tests/func/experiments/test_save.py
index 227b43a82..2427a00dd 100644
--- a/tests/func/experiments/test_save.py
+++ b/tests/func/experiments/test_save.py
@@ -49,7 +49,7 @@ def test_exp_save_overwrite_experiment(tmp_dir, dvc, scm, exp_stage):
     dvc.experiments.save(name="dummy", force=True)
 
 
-def test_exp_save_multiple(tmp_dir, dvc, scm, exp_stage):
+def test_exp_save_multiple(tmp_dir, dvc, scm):
     baseline = scm.get_rev()
     for i in range(2):
         name = f"exp-{i}"
@@ -76,3 +76,13 @@ def test_exp_save_after_commit(tmp_dir, dvc, scm, exp_stage):
     all_exps = dvc.experiments.ls(all_commits=True)
     assert all_exps[baseline] == ["exp-1"]
     assert all_exps[new_baseline] == ["exp-2"]
+
+
+def test_exp_save_with_staged_changes(tmp_dir, dvc, scm):
+    tmp_dir.gen({"new_file": "new_file"})
+    scm.add("new_file")
+
+    dvc.experiments.save(name="exp")
+
+    _, _, unstaged = scm.status()
+    assert "new_file" in unstaged
\ No newline at end of file