brick-graphviz.py

from graphviz import Digraph
import sys
from urllib.parse import quote
import re
import brickschema
from brickschema.namespaces import BRICK

classg = brickschema.Graph(load_brick=True)
typenodes = {}

colors = {
    BRICK.Location: "#ea9999",
    BRICK.Point: "#f0d890",
    BRICK.Equipment: "#b6d7a8",
}

G = Digraph("unified-model", engine='dot')
G.attr(ranksep=".3")
G.attr(fontsize="20")
G.attr(dpi="300")
# G.attr(width="600pt", height="600pt")

def get_name(node, graph):
    if node.startswith('ex:'):
        _, name = str(node).split(':')
        return name
    # return re.split(r'[/#:]', node)[-1]
    return graph.namespace_manager.normalizeUri(node)


def get_type_color(node):
    for (klass, color) in colors.items():
        if classg.query(f"ASK {{ <{node}> rdfs:subClassOf* <{klass}> }}").askAnswer:
            return color
    return "white"


def draw_node(n, graph, G):
    name = get_name(n, graph)
    if n.startswith('ex:'):
        G.node(quote(n), label=name, penwidth="4", fixedsize="false")
        return quote(n)
    else:
        color = get_type_color(n)
        nodename = quote(n)
        while nodename in str(G):
            nodename += "1"
        G.node(nodename, label=name, fillcolor=color, color=color, style="filled", fixedsize="false")
        return nodename


g = brickschema.Graph()
g.load_file(sys.argv[1])

# add all example nodes to graph
for (s, p, o) in g:
    s = draw_node(s, g, G)
    o = draw_node(o, g, G)

    p_name = get_name(p, g)
    if p_name == "rdf:type":
        G.edge(s, o, penwidth="2", arrowhead="diamond", style="dashed")
    else:
        G.edge(s, o, label=p_name, penwidth="2")


# G.attr(size="1200,1200")
# G.attr(fixedsize='true')
G.format = 'png'
G.render()