ffmap-backend/outputs/json/__init__.py
2014-07-07 23:27:21 +02:00

72 lines
2.1 KiB
Python

import json
__all__ = ["Exporter"]
class CustomJSONEncoder(json.JSONEncoder):
"""
JSON encoder that uses an object's __json__() method to convert it to
something JSON-compatible.
"""
def default(self, obj):
try:
return obj.__json__()
except AttributeError:
pass
return super().default(obj)
class Exporter:
def __init__(self, filepath="nodes.json"):
self.filepath = filepath
@staticmethod
def generate(nodedb):
indexes = {}
nodes = []
count = 0
for node in set(nodedb.values()):
nodes.append(node.export())
indexes[node.id] = count
count += 1
links = []
for node in set(nodedb.values()):
if "neighbors" in node:
links.extend(
{
"source": indexes[node.id],
"target": indexes[neighbor["neighbor"].id],
"quality": neighbor["metric"],
"type": "vpn" if neighbor["neighbor"]["vpn"] else None,
"id": "-".join((node.id, neighbor["neighbor"].id)),
} for neighbor in node["neighbors"]
)
if "clients" in node:
for client in node["clients"]:
if not client in indexes:
nodes.append({
"id": client,
})
indexes[client] = count
count += 1
links.append({
"source": indexes[node.id],
"target": indexes[client],
"quality": "TT",
"type": "client",
"id": "-".join((node.id, client)),
})
return {
"nodes": nodes,
"links": links,
}
def export(self, nodedb):
with open(self.filepath, "w") as nodes_json:
json.dump(
self.generate(nodedb),
nodes_json,
cls=CustomJSONEncoder
)