super.py

import json

from collections import defaultdict
from pathlib import Path


def load_from_json(path):
    with open(path, "r") as infile:
        return json.load(infile)


def dump_to_json(obj, path):
    with open(path, "w") as outfile:
        json.dump(obj, outfile, indent=4)


def extract_children(node_id, semantic_graph, super_node_ids):
    children = []
    for child in semantic_graph[node_id]:
        # if node isn't a super node, add it to list
        if child not in super_node_ids:
            children.append(child)

        # if child has children, recursively dive deeper again
        if child in semantic_graph:
            children += extract_children(child, semantic_graph, super_node_ids)

    return children


def get_super_nodes(semantic_graph, super_node_ids):
    super_mapping = {}
    for parent, children in semantic_graph.items():
        if parent in super_node_ids:
            # recursively grab all pgm nodes in the tree of nodes beneath it
            super_mapping[parent] = extract_children(parent, semantic_graph, super_node_ids)

    return super_mapping


def get_semantic_graph(path):
    semantic_graph_data = load_from_json(path)

    semantic_graph = defaultdict(list)

    for parent, child in semantic_graph_data:
        semantic_graph[parent].append(child)

    return semantic_graph


def validate_super_nodes(super_nodes_dict, super_node_ids):
    for super_node, children in super_nodes_dict.items():
        assert super_node in super_node_ids
        assert all([child_id not in super_node_ids for child_id in children])


def main():
    super_node_ids = {
        node["id"] for node in load_from_json(Path("nodes_all.json")) if node["label"] == "Super"
    }

    # load the semantic graph as a dict {parent: [children]}
    semantic_graph = get_semantic_graph(Path("semantic_graph.json"))
    
    # extract super nodes as dict {super_node_id: [all pgm nodes beneath it]}
    super_nodes = get_super_nodes(semantic_graph, super_node_ids)

    validate_super_nodes(super_nodes, super_node_ids)

    dump_to_json(super_nodes, Path("super_nodes.json"))


if __name__ == "__main__":
    main()