Merge pull request #17 from StructuralPython/features/add_flatten_function

Features/add flatten function

Author: connorferster

src/jsonchain/__init__.py

@@ -21,4 +21,6 @@
     filter_keys,
     merge_trees
 )
-from . import tables
+from . import tables
+
+from .tables import flatten_tree

src/jsonchain/tables.py

@@ -151,4 +151,38 @@ def create_tree_table(
             tree_leaves.update({header_row[idx]: row[idx]})
         tree_branch.update(tree_leaves)
     return tree_acc 
-    
+
+
+def flatten_tree(
+        tree: dict, 
+        level_labels: list[str | float | int], 
+        _current_level=0, 
+        _current_dict={}
+    ) -> list[dict]:
+    """
+    Returns a flattened list of dictionaries from the list.
+
+    tree: The nested dictionary to flatten. Must have the same number of
+        depth of across all branches.
+    level_labels: A list of strings to use as keys for each level.
+    current_level: The current depth of the recursion.
+    current_dict: The dictionary being built for the current path.
+    """
+    flattened_list = []
+
+    # Iterate over the key-value pairs of the current dictionary level
+    for key, value in tree.items():
+        new_dict = _current_dict.copy()
+        new_dict[level_labels[_current_level]] = key
+
+        # If the value is a dictionary and there are more levels to go, recurse
+        if isinstance(value, dict) and _current_level + 1 < len(level_labels):
+            flattened_list.extend(
+                flatten_tree(value, level_labels, _current_level + 1, new_dict)
+            )
+        else:
+            # Otherwise, we've reached a leaf node, so we update and append
+            new_dict.update(value)
+            flattened_list.append(new_dict)
+
+    return flattened_list

tests/test_jsonchain.py

@@ -1,4 +1,4 @@
-from jsonchain import load_json, dump_json, extract_keys
+from jsonchain import load_json, dump_json, extract_keys, flatten_tree
 import pathlib
 
 here = pathlib.Path.cwd()
@@ -31,4 +31,56 @@ def test_extract_keys():
         {"group": "aa"},
         {"group": "ab"},
     ]
-    
+    
+
+def test_flatten_tree():
+    tree = {
+        "a": {
+            "A": {
+                1: {
+                    "opta": 1,
+                    "optb": 2,
+                    "optc": 3,
+                },
+                2: {
+                    "opta": 10,
+                    "optb": 20,
+                    "optc": 30,
+                },
+            },
+            "B": {
+                1: {
+                    "opta": 21,
+                    "optb": 22,
+                    "optc": 23,
+                },
+                2: {
+                    "opta": 210,
+                    "optb": 220,
+                    "optc": 230,
+                }
+            },
+            "C": {
+                1: {
+                    "opta": 31,
+                    "optb": 32,
+                    "optc": 33,
+                },
+                2: {
+                    "opta": 310,
+                    "optb": 320,
+                    "optc": 330,
+                }
+            }
+        }
+    }
+    flat = [
+        {"member": "a", "force": "A", "case": 1, "opta": 1, "optb": 2, "optc": 3},
+        {"member": "a", "force": "A", "case": 2, "opta": 10, "optb": 20, "optc": 30},
+        {"member": "a", "force": "B", "case": 1, "opta": 21, "optb": 22, "optc": 23},
+        {"member": "a", "force": "B", "case": 2, "opta": 210, "optb": 220, "optc": 230},
+        {"member": "a", "force": "C", "case": 1, "opta": 31, "optb": 32, "optc": 33},
+        {"member": "a", "force": "C", "case": 2, "opta": 310, "optb": 320, "optc": 330},
+    ]
+    flattened = flatten_tree(tree, level_labels=['member', 'force', 'case'])
+    assert flattened == flat