add deep_set util function

Author: soxofaan

openeo/util.py

@@ -308,7 +308,7 @@ def __init__(self, key, keys):
 
 def deep_get(data: dict, *keys, default=_deep_get_default_undefined):
     """
-    Get "deep" value from nested dictionaries/lists/tuples
+    Get value deeply from nested dictionaries/lists/tuples
 
     :param data: nested data structure of dicts, lists, tuples
     :param keys: sequence of keys/indexes to traverse
@@ -329,6 +329,27 @@ def deep_get(data: dict, *keys, default=_deep_get_default_undefined):
     return data
 
 
+def deep_set(data: dict, *keys, value):
+    """
+    Set a value deeply in nested dictionary
+
+    :param data: nested data structure of dicts, lists, tuples
+    :param keys: sequence of keys/indexes to traverse
+    :param value: value to set
+    """
+    if len(keys) == 1:
+        data[keys[0]] = value
+    elif len(keys) > 1:
+        if isinstance(data, dict):
+            deep_set(data.setdefault(keys[0], {}), *keys[1:], value=value)
+        elif isinstance(data, (list, tuple)):
+            deep_set(data[keys[0]], *keys[1:], value=value)
+        else:
+            ValueError(data)
+    else:
+        raise ValueError("No keys given")
+
+
 def load_json(path: Union[Path, str]) -> dict:
     with Path(path).open("r", encoding="utf-8") as f:
         return json.load(f)

tests/test_util.py

@@ -9,7 +9,7 @@
 import pytest
 
 from openeo.util import first_not_none, get_temporal_extent, TimingLogger, ensure_list, ensure_dir, dict_no_none, \
-    deep_get, DeepKeyError, get_user_config_dir, get_user_data_dir, Rfc3339, rfc3339
+    deep_get, DeepKeyError, get_user_config_dir, get_user_data_dir, Rfc3339, rfc3339, deep_set
 
 
 def test_rfc3339_date():
@@ -356,6 +356,24 @@ def test_deep_get_mixed():
         deep_get(d, "bar", 2, 22, 222)
 
 
+@pytest.mark.parametrize(["init", "keys", "expected"], [
+    ({}, ("foo",), {"foo": 42}),
+    ({}, ("foo", "bar", "baz"), {"foo": {"bar": {"baz": 42}}}),
+    ({"foo": {"x": "y"}}, ("foo", "bar"), {"foo": {"x": "y", "bar": 42}}),
+    ({"foo": {"x": "y"}}, ("foo", "bar", "baz"), {"foo": {"x": "y", "bar": {"baz": 42}}}),
+    ({"foo": [1, 2, 3]}, ("foo", 1), {"foo": [1, 42, 3]}),
+    ({"foo": {1: "a", 2: "b"}}, ("foo", 1), {"foo": {1: 42, 2: "b"}}),
+    ({"foo": [{"x": 1}, {"x": 2}, {"x": 3}]}, ("foo", 1, "x"), {"foo": [{"x": 1}, {"x": 42}, {"x": 3}]}),
+    ({"foo": ({"x": 1}, {"x": 2}, {"x": 3})}, ("foo", 1, "x"), {"foo": ({"x": 1}, {"x": 42}, {"x": 3})}),
+    ({"foo": [{"x": {}}, {"x": {}}]}, ("foo", 1, "x", "bar"), {"foo": [{"x": {}}, {"x": {"bar": 42}}]}),
+    ({"foo": [{"x": {}}, {"x": {}}]}, ("foo", 1, "x", "y", "z"), {"foo": [{"x": {}}, {"x": {"y": {"z": 42}}}]}),
+])
+def test_deep_set_dict(init, keys, expected):
+    d = init
+    deep_set(d, *keys, value=42)
+    assert d == expected
+
+
 def test_get_user_config_dir():
     assert get_user_config_dir() == pathlib.Path(__file__).parent / "data/user_dirs/config/openeo-python-client"