lola_utils/LoLADraw.py

import subprocess

try:
    import pydot

    has_pydot = True
except ImportError:
    has_pydot = False

class LoLADraw:

    def __init__(self):
        print("Starting LoLADraw...")
        print("Please ensure that GraphViz is installed before use.")

    def get_net(self, filename):


        places = []
        markings = {}
        transitions = {}

        place_mode = False
        marking_mode = False
        transition_mode = False

        trans_consume_mode = True
        trans_name = ""
        trans_consume = {}
        trans_produce = {}

        with open(filename) as f:

            for line in f:
                if "PLACE" in line:
                    place_mode = True
                    continue
                elif "MARKING" in line:
                    place_mode = False
                    marking_mode = True
                    continue
                elif "TRANSITION" in line:
                    marking_mode = False
                    transition_mode = True

                line = line.strip()
                if not line:
                    continue

                print(line)
                if place_mode:
                    if "SAFE" in line:
                        continue
                    for place in line.split(','):
                        place = place.replace(";", "").strip()
                        if place:
                            places.append(place)
                elif marking_mode:
                    for marking in line.split(','):
                        if not marking:
                            continue
                        m = marking.split(":")
                        markings[m[0].strip()] = m[1].replace(";", "").strip()
                elif transition_mode:
                    if "TRANSITION" in line:

                        if trans_name:
                            transitions[trans_name] = (trans_consume, trans_produce)
                            trans_consume = {}
                            trans_produce = {}

                        trans_name = line.split(" ")[1]

                    elif trans_consume_mode:

                        if "CONSUME" in line:
                            continue
                        elif "PRODUCE" in line:
                            trans_consume_mode = False
                        else:
                            for consume in line.split(","):
                                if not consume:
                                    continue
                                place, marking = consume.split(":")
                                trans_consume[place.strip()] = marking.replace(";", "").strip()
                    elif not trans_consume_mode:
                        if ";" in line:
                            trans_consume_mode = True
                        for produce in line.split(","):
                            if not produce:
                                continue
                            place, marking = produce.split(":")
                            trans_produce[place.strip()] = marking.replace(";", "").strip()

            # do end-of-file
            if trans_name:
                transitions[trans_name] = (trans_consume, trans_produce)
        return places, markings, transitions

    def draw(self, filename):

        name = filename.replace(".lola", "")
        nodes = {}
        graph = pydot.Dot(name, graph_type='digraph', layout='neato')

        places, markings, transitions = self.get_net(filename)
        print("Places: " + str(places))
        print("Markings: " + str(markings))
        print("Transitions: " + str(transitions))

        for i, place in enumerate(places):
            v_attr = place
            if place in markings:
                v_attr += "\n" + markings[place]

            nodes[place] = pydot.Node(v_attr, style="filled", fillcolor='lightblue')
            graph.add_node(nodes[place])

        for t_name, _ in transitions.items():
            v_attr = t_name

            nodes[t_name] = pydot.Node(v_attr, style="filled")
            graph.add_node(nodes[t_name])

        for t_name, conpro in transitions.items():
            c, p = conpro

            for c_name, weight in c.items():
                graph.add_edge(pydot.Edge(nodes[c_name], nodes[t_name]))#, color=color))

            for p_name, weight in p.items():
                graph.add_edge(pydot.Edge(nodes[t_name], nodes[p_name], label=weight))
                # , color=color))


        dot_filename = filename.replace(".lola", ".dot")
        graph.write(dot_filename, prog='dot')

        command = "dot -Tsvg " + dot_filename + " -o " + filename.replace(".lola", ".svg")
        subprocess.call(command, shell=True)

        # command = "rm " + dot_filename
        # subprocess.call(command, shell=True)


if __name__ == "__main__":

    if not has_pydot:
        print("ERROR: Install pydot first!")
        exit(1)

    filename = "dining_philo.lola"
    ld = LoLADraw()
    ld.draw(filename)