Merge nested env keys when using import_from (#1034)

src/ploomber/env/envdict.py

@@ -131,7 +131,7 @@ def __init__(self, source, path_to_here=None, defaults=None):
             import_from = _get_import_from(raw_data, path_to_here)
 
             if import_from:
-                raw_data = {**import_from, **raw_data}
+                raw_data = deep_merge(import_from, raw_data)
 
             # check raw data is ok
             validate.raw_data_keys(raw_data)
@@ -472,3 +472,14 @@ def find_tags_in_dict(d):
         tags = tags | util.get_tags_in_str(k)
 
     return tags
+
+
+def deep_merge(a: dict, b: dict) -> dict:
+    result = deepcopy(a)
+    for bk, bv in b.items():
+        av = result.get(bk)
+        if isinstance(av, dict) and isinstance(bv, dict):
+            result[bk] = deep_merge(av, bv)
+        else:
+            result[bk] = deepcopy(bv)
+    return result

tests/env/test_env.py

@@ -936,20 +936,28 @@ def test_import_from_another_file(tmp_directory):
     Path('base.yaml').write_text("""
 key: value
 base: 42
+nested:
+  key_1: value_1
+  key_2: value_2
 """)
 
     env = EnvDict(
         {
             'meta': {
                 'import_from': 'base.yaml'
             },
-            'key': 'new_value'
+            'key': 'new_value',
+            'nested': {
+                'key_1': 'new_value_1',
+            },
         },
         path_to_here=os.getcwd(),
     )
 
     assert env['key'] == 'new_value'
     assert env['base'] == 42
+    assert env['nested']['key_1'] == 'new_value_1'
+    assert env['nested']['key_2'] == 'value_2'
     # this section is for config, and should not be visible
     assert 'meta' not in env