File size: 2,626 Bytes
4dd7afe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | 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, [])
|