Commit 818fc0f6 authored by Stanislav Bohm's avatar Stanislav Bohm
Browse files

ENH: Plan to .dot file

parent e2f7a45a
class Arc(object):
def __init__(self, node, data, color, style):
self.node = node = data
self.color = color = style
class Node(object):
color = None
fontcolor = None
fillcolor = None
label = ""
html_label = None
shape = None
def __init__(self, key):
self.key = key
self.arcs = []
def add_arc(self, node, data=None, color=None, style=None):
self.arcs.append(Arc(node, data, color, style))
def merge_arcs(self, merge_fn):
if len(self.arcs) < 2:
node_to_arcs = {}
for arc in self.arcs[:]:
a = node_to_arcs.get(arc.node)
if a is None:
node_to_arcs[arc.node] = arc
self.arcs.remove(arc) = merge_fn(,
def __repr__(self):
return "<Node {}>".format(self.key)
class Graph(object):
def __init__(self):
self.nodes = {}
def size(self):
return len(self.nodes)
def node_check(self, key):
node = self.nodes.get(key)
if node is not None:
return (node, True)
node = Node(key)
self.nodes[key] = node
return (node, False)
def node(self, key):
node = self.nodes.get(key)
if node is not None:
return node
node = Node(key)
self.nodes[key] = node
return node
def show(self):
def write(self, filename):
dot = self.make_dot("G")
with open(filename, "w") as f:
def make_dot(self, name):
stream = ["digraph " + name + " {\n"]
stream.append("node [shape=box];")
for node in self.nodes.values():
extra = ""
if node.fontcolor is not None:
extra += " fontcolor=\"{}\"".format(node.fontcolor)
if node.shape is not None:
extra += " shape={}".format(node.shape)
if node.color is not None:
extra += " color=\"{}\"".format(node.color)
if node.fillcolor is not None:
extra += " style=filled fillcolor=\"{}\""\
if node.html_label is not None:
label = "<" + node.html_label + ">"
label = "\"" + node.label + "\""
stream.append("v{} [label={}{}]\n".format(
id(node), label, extra))
for arc in node.arcs:
extra = ""
if is not None:
extra = "label=\"{}\"".format(
if arc.color is not None:
extra += " color=\"{}\"".format(arc.color)
if is not None:
extra += " style=\"{}\"".format(
stream.append("v{} -> v{} [{}]\n".format(
id(node), id(arc.node), extra))
return "".join(stream)
def merge_arcs(self, merge_fn):
for node in self.nodes.values():
def run_xdot(dot):
import subprocess
import tempfile
with tempfile.NamedTemporaryFile() as f:
import loomplan_pb2
import loomrun_pb2
import gv
import struct
......@@ -173,3 +175,12 @@ class Plan(object):
t = msg.tasks.add()
task.set_message(t, task_types)
return msg
def write_dot(self, filename):
graph = gv.Graph()
for task in self.tasks:
node = graph.node(
node.label = "{}\n{}".format(str(, task.task_type)
for inp in task.inputs:
......@@ -38,6 +38,7 @@ def test_cv_iris(loom_env):
task.map_file_in(model, "model")
# p.write_dot("")
results = loom_env.submit(p, predict)
assert len(results) == CHUNKS
for line in results:
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment