From 7444696dcd114bac4a80942b0eca34ab87bab5ad Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Tue, 3 Aug 2021 01:30:04 +0545
Subject: [PATCH] Fix merging of params when empty (#6378)

* Fix merging of params when empty

Recent release of python-benedict==0.24.1 started ignoring empty dictionaries
during merge. The way we expected to work was questionable and also how
benedict itself tries to handle the references to empty dictionary is
also questionable.

I consider this to be a proper way to handle this even if this gets fixed
on the benedict side.

* Revert "setup: limit python-benedict version (#6377)"

This reverts commit 15de826266c913a967a2913293f59c65b333fdf1.

* add comment for why passing through benedict
---
 dvc/repo/experiments/__init__.py     |  5 +-
 dvc/utils/collections.py             | 14 +++++
 setup.py                             |  2 +-
 tests/unit/utils/test_collections.py | 76 +++++++++++++++++++++++++++-
 4 files changed, 92 insertions(+), 5 deletions(-)

diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index 5a560c14c..06f326d9e 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/dvc/repo/experiments/__init__.py
@@ -352,8 +352,7 @@ class Experiments:
 
     def _update_params(self, params: dict):
         """Update experiment params files with the specified values."""
-        from benedict import benedict
-
+        from dvc.utils.collections import merge_params
         from dvc.utils.serialize import MODIFIERS
 
         logger.debug("Using experiment params '%s'", params)
@@ -363,7 +362,7 @@ class Experiments:
             suffix = path.suffix.lower()
             modify_data = MODIFIERS[suffix]
             with modify_data(path, fs=self.repo.fs) as data:
-                benedict(data).merge(params[params_fname], overwrite=True)
+                merge_params(data, params[params_fname])
 
         # Force params file changes to be staged in git
         # Otherwise in certain situations the changes to params file may be
diff --git a/dvc/utils/collections.py b/dvc/utils/collections.py
index 70ecf56e2..15363180d 100644
--- a/dvc/utils/collections.py
+++ b/dvc/utils/collections.py
@@ -80,6 +80,20 @@ def chunk_dict(d: Dict[_KT, _VT], size: int = 1) -> List[Dict[_KT, _VT]]:
     return [{key: d[key] for key in chunk} for chunk in chunks(size, d)]
 
 
+def merge_params(src: Dict, to_update: Dict) -> Dict:
+    """Recursively merges params with benedict's syntax support in-place."""
+    from benedict import benedict
+
+    if src:
+        benedict(src).merge(to_update, overwrite=True)
+    else:
+        # benedict has issues keeping references to an empty dictionary
+        # see: https://github.com/iterative/dvc/issues/6374.
+        # Also, passing to_update through benedict to expand the syntax.
+        src.update(benedict(to_update))
+    return src
+
+
 class _NamespacedDict(dict):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
diff --git a/setup.py b/setup.py
index ddf39a746..67f17a8f3 100644
--- a/setup.py
+++ b/setup.py
@@ -85,7 +85,7 @@ install_requires = [
     "shtab>=1.3.4,<2",
     "rich>=10.0.0",
     "dictdiffer>=0.8.1",
-    "python-benedict>=0.21.1,<0.24.1",
+    "python-benedict>=0.21.1",
     "pyparsing==2.4.7",
     "typing_extensions>=3.7.4",
     "fsspec>=2021.7.0",
diff --git a/tests/unit/utils/test_collections.py b/tests/unit/utils/test_collections.py
index 9de4b1917..f7b02d1c1 100644
--- a/tests/unit/utils/test_collections.py
+++ b/tests/unit/utils/test_collections.py
@@ -1,7 +1,13 @@
 # pylint: disable=unidiomatic-typecheck
+import pytest
 from mock import create_autospec
 
-from dvc.utils.collections import apply_diff, chunk_dict, validate
+from dvc.utils.collections import (
+    apply_diff,
+    chunk_dict,
+    merge_params,
+    validate,
+)
 
 
 class MyDict(dict):
@@ -135,3 +141,71 @@ def test_post_validate_decorator(mocker):
     result = func()
     test_func.assert_called_once()
     assert result == [1, 2]
+
+
+@pytest.mark.parametrize(
+    "changes, expected",
+    [
+        [{"foo": "baz"}, {"foo": "baz", "goo": {"bag": 3}, "lorem": False}],
+        [
+            {"foo": "baz", "goo": "bar"},
+            {"foo": "baz", "goo": "bar", "lorem": False},
+        ],
+        [
+            {"goo.bag": 4},
+            {"foo": {"bar": 1, "baz": 2}, "goo": {"bag": 4}, "lorem": False},
+        ],
+        [
+            {"foo[0]": "bar"},
+            {
+                "foo": {"bar": 1, "baz": 2, 0: "bar"},
+                "goo": {"bag": 3},
+                "lorem": False,
+            },
+        ],
+        [
+            {"foo[1].baz": 3},
+            {
+                "foo": {"bar": 1, "baz": 2, 1: {"baz": 3}},
+                "goo": {"bag": 3},
+                "lorem": False,
+            },
+        ],
+        [
+            {"foo[1]": ["baz", "goo"]},
+            {
+                "foo": {"bar": 1, "baz": 2, 1: ["baz", "goo"]},
+                "goo": {"bag": 3},
+                "lorem": False,
+            },
+        ],
+        [
+            {"lorem.ipsum": 3},
+            {
+                "foo": {"bar": 1, "baz": 2},
+                "goo": {"bag": 3},
+                "lorem": {"ipsum": 3},
+            },
+        ],
+        [{}, {"foo": {"bar": 1, "baz": 2}, "goo": {"bag": 3}, "lorem": False}],
+    ],
+)
+def test_merge_params(changes, expected):
+    params = {"foo": {"bar": 1, "baz": 2}, "goo": {"bag": 3}, "lorem": False}
+    merged = merge_params(params, changes)
+    assert merged == expected == params
+    assert params is merged  # references should be preserved
+
+
+@pytest.mark.parametrize(
+    "changes, expected",
+    [
+        [{"foo": "baz"}, {"foo": "baz"}],
+        [{"foo": "baz", "goo": "bar"}, {"foo": "baz", "goo": "bar"}],
+    ],
+)
+def test_merge_params_on_empty_src(changes, expected):
+    params = {}
+    merged = merge_params(params, changes)
+    assert merged == expected == params
+    assert params is merged  # references should be preserved