dump_janusgraph.py

#!/usr/bin/env nix-shell
#!nix-shell -p jq -i python3 -p python3Packages.requests -p graphviz

# NOTE: Edit the dict below to change custom highlight colors for the labels of the nodes/vertices.
bgcolors = { "_key": "pink", "thing": "yellow" , "action": "green"}

import requests
import json
import subprocess


def get_gremlin(gremlin):
    response = requests.post('http://localhost:8182', json= {'gremlin': gremlin})
    if response.status_code != 200:
        raise Exception("Not 200 status;", response)
    return response.json()["result"]["data"]

class Graph:
    def __init__(self):
        self.g = "digraph G {\n"
    def add(self, x):
        self.g += "  " + x + "\n"
    def dot(self):
        return self.g + "}\n"

def html_dot_escape(input):
    return input.replace("&", "&amp;").replace('"',"&quot;").replace("<", "&lt;").replace(">", "&gt;")

def dump_to_graphviz():
    vertices = get_gremlin("g.V()")
    edges    = get_gremlin("g.E()")
    g = Graph()

    for vertex in vertices:
        label  = '<TABLE><TR><TD BGCOLOR="blue">vertex</TD><TD BGCOLOR="{}">{}</TD></TR>'.format(bgcolors.get(vertex["label"], "white"), vertex["label"])
        label += '<TR><TD>ID</TD><TD>{}</TD></TR>'.format(vertex["id"])

        items = list(vertex.get("properties", {}).items())
        items.sort(key=lambda x: x[0])
        for prop_name, prop_values in items:
            for prop_value in prop_values:
                label += "<TR><TD>{}</TD><TD>{}</TD></TR>".format(prop_name, html_dot_escape(str(prop_value["value"])))

        label += '</TABLE>'
        g.add('v{}[shape=none,label=<{}>]'.format(vertex["id"], label))


    for edge in edges:
        label  = '<TABLE><TR><TD BGCOLOR="green">edge</TD><TD BGCOLOR="{}">{}</TD></TR>'.format(bgcolors.get(edge["label"], "white"), edge["label"])
        label += '<TR><TD>ID</TD><TD>{}</TD></TR>'.format(edge["id"])

        items = list(edge.get("properties", {}).items())
        items.sort(key=lambda x: x[0])
        for prop_name, prop_value in items:
            #print("'{}': '{}'".format(prop_name, prop_value))
            label += "<TR><TD>{}</TD><TD>{}</TD></TR>".format(prop_name, html_dot_escape(str(prop_value)))

        label += '</TABLE>'
        edge_id = edge["id"].replace("-","_")
        g.add('e{}[shape=none,label=<{}>]'.format(edge_id, label))

        g.add("v{} -> e{}; e{}->v{}".format(edge["outV"], edge_id, edge_id, edge["inV"]))

    return g.dot()


if __name__ == "__main__":
    graphviz = dump_to_graphviz()
    print(graphviz)
    subprocess.check_output(["dot", "-Tx11"], input=bytes(graphviz, encoding="utf-8"))