From 831e85682e341d47f341ea27decbbf79f534783c Mon Sep 17 00:00:00 2001
From: Sami Jawhar <sami@thecybermonk.com>
Date: Fri, 23 Oct 2020 00:35:17 -0500
Subject: [PATCH] Experiments: allow setting of list parameters (#4769)

* Allow setting of list parameteres

Check if dict is Mapping or list before using .get()

* Added test for dvc exp run with list paramater

* Add more tests for string parameter modification
---
 dvc/repo/experiments/__init__.py           |  8 +++++-
 tests/func/experiments/test_experiments.py | 29 ++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index e2468704c..55e6c2607 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/dvc/repo/experiments/__init__.py
@@ -318,8 +318,14 @@ class Experiments:
         # recursive dict update
         def _update(dict_, other):
             for key, value in other.items():
+                if isinstance(dict_, list) and key.isdigit():
+                    key = int(key)
                 if isinstance(value, Mapping):
-                    dict_[key] = _update(dict_.get(key, {}), value)
+                    if isinstance(dict_, list):
+                        fallback_value = dict_[key]
+                    else:
+                        fallback_value = dict_.get(key, {})
+                    dict_[key] = _update(fallback_value, value)
                 else:
                     dict_[key] = value
             return dict_
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index afe970481..74c3d2c7b 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -93,6 +93,35 @@ def test_update_with_pull(tmp_dir, scm, dvc, mocker):
         assert exp_scm.has_rev(rev)
 
 
+@pytest.mark.parametrize(
+    "change, expected",
+    [
+        ["foo.1.baz=3", "foo: [bar: 1, baz: 3]"],
+        ["foo.0=bar", "foo: [bar, baz: 2]"],
+        ["foo.1=- baz\n- goo", "foo: [bar: 1, [baz, goo]]"],
+    ],
+)
+def test_modify_list_param(tmp_dir, scm, dvc, mocker, change, expected):
+    tmp_dir.gen("copy.py", COPY_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: [bar: 1, baz: 2]")
+    stage = dvc.run(
+        cmd="python copy.py params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        name="copy-file",
+    )
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("init")
+
+    new_mock = mocker.spy(dvc.experiments, "new")
+    dvc.experiments.run(stage.addressing, params=[change])
+
+    new_mock.assert_called_once()
+    assert (
+        tmp_dir / ".dvc" / "experiments" / "metrics.yaml"
+    ).read_text().strip() == expected
+
+
 def test_checkout(tmp_dir, scm, dvc):
     tmp_dir.gen("copy.py", COPY_SCRIPT)
     tmp_dir.gen("params.yaml", "foo: 1")