summaryrefslogtreecommitdiff
path: root/tools/state_graph.py
blob: f6a2057dd30a360c14f29e6a62001a261be94040 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#!/usr/bin/env python

import os
import sys

top_dir = os.path.abspath(os.path.join(os.path.dirname(__file__),
                                       os.pardir))
sys.path.insert(0, top_dir)

import optparse
import subprocess
import tempfile

from taskflow import states
from taskflow.types import graph as gr


def mini_exec(cmd, ok_codes=(0,)):
    stdout = subprocess.PIPE
    stderr = subprocess.PIPE
    proc = subprocess.Popen(cmd, stdout=stdout, stderr=stderr, stdin=None)
    (stdout, stderr) = proc.communicate()
    rc = proc.returncode
    if rc not in ok_codes:
        raise RuntimeError("Could not run %s [%s]\nStderr: %s"
                           % (cmd, rc, stderr))
    return (stdout, stderr)


def make_svg(graph, output_filename, output_format):
    # NOTE(harlowja): requires pydot!
    gdot = graph.export_to_dot()
    if output_format == 'dot':
        output = gdot
    elif output_format in ('svg', 'svgz', 'png'):
        with tempfile.NamedTemporaryFile(suffix=".dot") as fh:
            fh.write(gdot)
            fh.flush()
            cmd = ['dot', '-T%s' % output_format, fh.name]
            output, _stderr = mini_exec(cmd)
    else:
        raise ValueError('Unknown format: %s' % output_filename)
    with open(output_filename, "wb") as fh:
        fh.write(output)


def main():
    parser = optparse.OptionParser()
    parser.add_option("-f", "--file", dest="filename",
                      help="write svg to FILE", metavar="FILE")
    parser.add_option("-t", "--tasks", dest="tasks",
                      action='store_true',
                      help="use task state transitions",
                      default=False)
    parser.add_option("-T", "--format", dest="format",
                      help="output in given format",
                      default='svg')

    (options, args) = parser.parse_args()
    if options.filename is None:
        options.filename = 'states.%s' % options.format

    g = gr.DiGraph(name="State transitions")
    if not options.tasks:
        source = states._ALLOWED_FLOW_TRANSITIONS
    else:
        source = states._ALLOWED_TASK_TRANSITIONS
    for (u, v) in source:
        if not g.has_node(u):
            g.add_node(u)
        if not g.has_node(v):
            g.add_node(v)
        g.add_edge(u, v)
    make_svg(g, options.filename, options.format)
    print("Created %s at '%s'" % (options.format, options.filename))


if __name__ == '__main__':
    main()