id_mapper.py

from typing import Dict, List, Tuple
import json

class UnifiedIdMapper:
    def __init__(self, nodes: Dict[int, str], edges: Dict[int, str]) -> None:

        # since all key in JSON are str, convert them to int
        nodes = {int(k): v for k, v in nodes.items()}
        edges = {int(k): v for k, v in edges.items()}

        self.nodes = nodes
        self.edges = edges

        node_mapping = {old_id: new_id for new_id, old_id in enumerate(sorted(self.nodes.keys()))}
        edge_mapping = {old_id: new_id for new_id, old_id in enumerate(sorted(edges.keys()))}
        shift = len(nodes)

        self.old_to_new: Dict[int, Tuple[int, bool]] = {
            **{old_id: (new_id, False) for old_id, new_id in node_mapping.items()},
            **{old_id: (new_id + shift, True) for old_id, new_id in edge_mapping.items()},
        }
        # reverse mapping: new_id -> (old_id, is_edge)
        self.new_to_old: Dict[int, Tuple[int, bool]] = {
            new_id: (old_id, is_edge)
            for old_id, (new_id, is_edge) in self.old_to_new.items()
        }

        # Label maps
        self.old_id_to_label: Dict[int, str] = {**nodes, **edges}
        self.new_id_to_label: Dict[int, str] = {
            new_id: self.old_id_to_label[old_id] for old_id, (new_id, _) in self.old_to_new.items()
        }

        self.label_to_old_ids: Dict[str, List[Tuple[int, bool]]] = {}
        self.label_to_new_ids: Dict[str, List[Tuple[int, bool]]] = {}
        for old_id, (new_id, is_edge) in self.old_to_new.items():
            label = self.old_id_to_label.get(old_id)
            if label is None:
                continue
            self.label_to_old_ids.setdefault(label, []).append((old_id, is_edge))
            self.label_to_new_ids.setdefault(label, []).append((new_id, is_edge))

    @classmethod
    def from_file(cls, mapper_path: str):
        with open(mapper_path, "r") as f:
            data = json.load(f)
            return cls(data['nodes'], data['edges'])

    def map_old_id(self, old_id: int) -> Tuple[int, bool]:
        return self.old_to_new[old_id]

    def map_new_id(self, new_id: int) -> Tuple[int, bool]:
        return self.new_to_old[new_id]

    def label_from_old_id(self, old_id: int) -> str:
        return self.old_id_to_label[old_id]

    def label_from_new_id(self, new_id: int) -> str:
        return self.new_id_to_label[new_id]

    def old_ids_from_label(self, label: str) -> List[Tuple[int, bool]]:
        return self.label_to_old_ids.get(label, [])

    def new_ids_from_label(self, label: str) -> List[Tuple[int, bool]]:
        return self.label_to_new_ids.get(label, [])