Adding array support for unflatten_dict

src/fides/api/util/saas_util.py

@@ -3,7 +3,6 @@
 import json
 import re
 from collections import defaultdict
-from functools import reduce
 from typing import Any, Dict, List, Optional, Set, Tuple
 
 import yaml
@@ -222,11 +221,18 @@ def unflatten_dict(flat_dict: Dict[str, Any], separator: str = ".") -> Dict[str,
                 "'unflatten_dict' expects a flattened dictionary as input."
             )
         keys = path.split(separator)
-        target = reduce(
-            lambda current, key: current.setdefault(key, {}),
-            keys[:-1],
-            output,
-        )
+        target = output
+        for i, current_key in enumerate(keys[:-1]):
+            next_key = keys[i + 1]
+            if next_key.isdigit():
+                target = target.setdefault(current_key, [])
+            else:
+                if isinstance(target, dict):
+                    target = target.setdefault(current_key, {})
+                elif isinstance(target, list):
+                    while len(target) <= int(current_key):
+                        target.append({})
+                    target = target[int(current_key)]
         try:
             target[keys[-1]] = value
         except TypeError as exc:

tests/ops/util/test_saas_util.py

@@ -329,50 +329,63 @@ def test_regex_is_not_greedy(self):
         )
 
 
-def test_unflatten_dict():
-    # empty dictionary
-    assert unflatten_dict({}) == {}
-
-    # empty dictionary value
-    assert unflatten_dict({"A": {}}) == {"A": {}}
+class TestUnflattenDict:
+    def test_empty_dict(self):
+        assert unflatten_dict({}) == {}
+
+    def test_empty_dict_value(self):
+        assert unflatten_dict({"A": {}}) == {"A": {}}
+
+    def test_unflattened_dict(self):
+        assert unflatten_dict({"A": "1"}) == {"A": "1"}
+
+    def test_same_level(self):
+        assert unflatten_dict({"A.B": "1", "A.C": "2"}) == {"A": {"B": "1", "C": "2"}}
+
+    def test_mixed_levels(self):
+        assert unflatten_dict(
+            {
+                "A": "1",
+                "B.C": "2",
+                "B.D": "3",
+            }
+        ) == {
+            "A": "1",
+            "B": {"C": "2", "D": "3"},
+        }
 
-    # unflattened dictionary
-    assert unflatten_dict({"A": "1"}) == {"A": "1"}
+    def test_long_path(self):
+        assert unflatten_dict({"A.B.C.D.E.F.G": "1"}) == {
+            "A": {"B": {"C": {"D": {"E": {"F": {"G": "1"}}}}}}
+        }
 
-    # same level
-    assert unflatten_dict({"A.B": "1", "A.C": "2"}) == {"A": {"B": "1", "C": "2"}}
+    def test_single_item_array(self):
+        assert unflatten_dict({"A.0.B": "C"}) == {"A": [{"B": "C"}]}
 
-    # mixed levels
-    assert unflatten_dict(
-        {
-            "A": "1",
-            "B.C": "2",
-            "B.D": "3",
+    def test_multi_item_array(self):
+        assert unflatten_dict({"A.0.B": "C", "A.1.D": "E"}) == {
+            "A": [{"B": "C"}, {"D": "E"}]
         }
-    ) == {
-        "A": "1",
-        "B": {"C": "2", "D": "3"},
-    }
-
-    # long path
-    assert unflatten_dict({"A.B.C.D.E.F.G": "1"}) == {
-        "A": {"B": {"C": {"D": {"E": {"F": {"G": "1"}}}}}}
-    }
-
-    # incoming values should overwrite existing values
-    assert unflatten_dict({"A.B": 1, "A.B": 2}) == {"A": {"B": 2}}
-
-    # conflicting types
-    with pytest.raises(FidesopsException):
-        unflatten_dict({"A.B": 1, "A": 2, "A.C": 3})
-
-    # data passed in is not completely flattened
-    with pytest.raises(FidesopsException):
-        unflatten_dict({"A.B.C": 1, "A": {"B.D": 2}})
-
-    # unflatten_dict shouldn't be called with a None separator
-    with pytest.raises(IndexError):
-        unflatten_dict({"": "1"}, separator=None)
+
+    def test_multi_value_array(self):
+        assert unflatten_dict(
+            {"A.0.B": "C", "A.0.D": "E", "A.1.F": "G", "A.1.H": "I"}
+        ) == {"A": [{"B": "C", "D": "E"}, {"F": "G", "H": "I"}]}
+
+    def test_overwrite_existing_values(self):
+        assert unflatten_dict({"A.B": 1, "A.B": 2}) == {"A": {"B": 2}}
+
+    def test_conflicting_types(self):
+        with pytest.raises(FidesopsException):
+            unflatten_dict({"A.B": 1, "A": 2, "A.C": 3})
+
+    def test_data_not_completely_flattened(self):
+        with pytest.raises(FidesopsException):
+            unflatten_dict({"A.B.C": 1, "A": {"B.D": 2}})
+
+    def test_none_separator(self):
+        with pytest.raises(IndexError):
+            unflatten_dict({"": "1"}, separator=None)
 
 
 def test_replace_version():