Add create_hierarchy_dict and create_hierarchy_tree functions

isco.py

@@ -1,8 +1,9 @@
 # filename: isco.py
+from collections import defaultdict
 import csv
 
 
-def create_hierarchy(filename):
+def create_hierarchy_dict(filename):
     """
     Creates a dictionary where keys are nodes and values are sets of parent nodes representing the hierarchy of the ISCO-08 codes from the "unit" column of the isco_structure.csv file.
 
@@ -34,6 +35,37 @@ def create_hierarchy(filename):
     return isco_hierarchy
 
 
+def create_hierarchy_tree(hierarchy_dict: dict) -> tuple:
+    """
+    Builds the hierarchy tree and a mapping from name to ISCO code.
+
+    Args:
+    - hierarchy_json: A dictionary representing the hierarchical structure.
+
+    Returns:
+    - tree: A dictionary representing the hierarchical structure.
+    - code_to_node: A dictionary mapping from ISCO code to node name.
+    """
+
+    tree = defaultdict(lambda: {"children": [], "parent": None})
+    code_to_node = {}
+
+    def add_node(parent_code, node):
+        code = node["name"].split("=")[0].strip()
+        code_to_node[code] = node["name"]
+        tree[code]["parent"] = parent_code
+        if parent_code:
+            tree[parent_code]["children"].append(code)
+        for child in node.get("children", []):
+            add_node(code, child)
+
+    add_node(None, hierarchy_dict)  # Root node has no parent
+    return tree, code_to_node
+
+
 # Example usage:
-# hierarchy = create_hierarchy("ISCO_structure.csv")
+# hierarchy_dict = create_hierarchy("ISCO_structure.csv")
+# tree, code_to_node = create_hierarchy_tree(hierarchy_dict)
 # print(hierarchy)
+# print(code_to_node)
+# print(tree)