054ca2a6e2
Just directly use the built machine in the action engine and avoid having another layer of abstraction that does not provide that much value. This makes the code cleaner, and more easy to understand (and so-on). Change-Id: Iae1279098112254338258c1941c15889f1ad1a79
215 lines
7.5 KiB
Python
Executable File
215 lines
7.5 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
# Copyright (C) 2014 Yahoo! Inc. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License"); you may
|
|
# not use this file except in compliance with the License. You may obtain
|
|
# a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
# License for the specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
import mock
|
|
|
|
import optparse
|
|
import os
|
|
import sys
|
|
|
|
top_dir = os.path.abspath(os.path.join(os.path.dirname(__file__),
|
|
os.pardir))
|
|
sys.path.insert(0, top_dir)
|
|
|
|
# To get this installed you have to do the following:
|
|
#
|
|
# $ pip install pydot2
|
|
import pydot
|
|
|
|
from automaton import machines
|
|
|
|
from taskflow.engines.action_engine import builder
|
|
from taskflow.engines.worker_based import protocol
|
|
from taskflow import states
|
|
|
|
|
|
# This is just needed to get at the machine object (we will not
|
|
# actually be running it...).
|
|
class DummyRuntime(object):
|
|
def __init__(self):
|
|
self.analyzer = mock.MagicMock()
|
|
self.completer = mock.MagicMock()
|
|
self.scheduler = mock.MagicMock()
|
|
self.storage = mock.MagicMock()
|
|
|
|
|
|
def clean_event(name):
|
|
name = name.replace("_", " ")
|
|
name = name.strip()
|
|
return name
|
|
|
|
|
|
def make_machine(start_state, transitions):
|
|
machine = machines.FiniteMachine()
|
|
machine.add_state(start_state)
|
|
for (start_state, end_state) in transitions:
|
|
if start_state not in machine:
|
|
machine.add_state(start_state)
|
|
if end_state not in machine:
|
|
machine.add_state(end_state)
|
|
# Make a fake event (not used anyway)...
|
|
event = "on_%s" % (end_state)
|
|
machine.add_transition(start_state, end_state, event.lower())
|
|
machine.default_start_state = start_state
|
|
return machine
|
|
|
|
|
|
def map_color(internal_states, state):
|
|
if state in internal_states:
|
|
return 'blue'
|
|
if state in (states.FAILURE, states.REVERT_FAILURE):
|
|
return 'red'
|
|
if state == states.REVERTED:
|
|
return 'darkorange'
|
|
if state in (states.SUCCESS, states.COMPLETE):
|
|
return 'green'
|
|
return None
|
|
|
|
|
|
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("-r", "--retries", dest="retries",
|
|
action='store_true',
|
|
help="use retry state transitions",
|
|
default=False)
|
|
parser.add_option("-e", "--engines", dest="engines",
|
|
action='store_true',
|
|
help="use engine state transitions",
|
|
default=False)
|
|
parser.add_option("-w", "--wbe-requests", dest="wbe_requests",
|
|
action='store_true',
|
|
help="use wbe request transitions",
|
|
default=False)
|
|
parser.add_option("-j", "--jobs", dest="jobs",
|
|
action='store_true',
|
|
help="use job 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
|
|
|
|
types = [
|
|
options.engines,
|
|
options.retries,
|
|
options.tasks,
|
|
options.wbe_requests,
|
|
options.jobs,
|
|
]
|
|
if sum([int(i) for i in types]) > 1:
|
|
parser.error("Only one of task/retry/engines/wbe requests/jobs"
|
|
" may be specified.")
|
|
|
|
internal_states = list()
|
|
ordering = 'in'
|
|
if options.tasks:
|
|
source_type = "Tasks"
|
|
source = make_machine(states.PENDING,
|
|
list(states._ALLOWED_TASK_TRANSITIONS))
|
|
elif options.retries:
|
|
source_type = "Retries"
|
|
source = make_machine(states.PENDING,
|
|
list(states._ALLOWED_RETRY_TRANSITIONS))
|
|
elif options.engines:
|
|
source_type = "Engines"
|
|
b = builder.MachineBuilder(DummyRuntime(), mock.MagicMock())
|
|
source, memory = b.build()
|
|
internal_states.extend(builder.META_STATES)
|
|
ordering = 'out'
|
|
elif options.wbe_requests:
|
|
source_type = "WBE requests"
|
|
source = make_machine(protocol.WAITING,
|
|
list(protocol._ALLOWED_TRANSITIONS))
|
|
elif options.jobs:
|
|
source_type = "Jobs"
|
|
source = make_machine(states.UNCLAIMED,
|
|
list(states._ALLOWED_JOB_TRANSITIONS))
|
|
else:
|
|
source_type = "Flow"
|
|
source = make_machine(states.PENDING,
|
|
list(states._ALLOWED_FLOW_TRANSITIONS))
|
|
|
|
graph_name = "%s states" % source_type
|
|
g = pydot.Dot(graph_name=graph_name, rankdir='LR',
|
|
nodesep='0.25', overlap='false',
|
|
ranksep="0.5", size="11x8.5",
|
|
splines='true', ordering=ordering)
|
|
node_attrs = {
|
|
'fontsize': '11',
|
|
}
|
|
nodes = {}
|
|
for (start_state, on_event, end_state) in source:
|
|
on_event = clean_event(on_event)
|
|
if start_state not in nodes:
|
|
start_node_attrs = node_attrs.copy()
|
|
text_color = map_color(internal_states, start_state)
|
|
if text_color:
|
|
start_node_attrs['fontcolor'] = text_color
|
|
nodes[start_state] = pydot.Node(start_state, **start_node_attrs)
|
|
g.add_node(nodes[start_state])
|
|
if end_state not in nodes:
|
|
end_node_attrs = node_attrs.copy()
|
|
text_color = map_color(internal_states, end_state)
|
|
if text_color:
|
|
end_node_attrs['fontcolor'] = text_color
|
|
nodes[end_state] = pydot.Node(end_state, **end_node_attrs)
|
|
g.add_node(nodes[end_state])
|
|
if options.engines:
|
|
edge_attrs = {
|
|
'label': on_event,
|
|
}
|
|
if 'reverted' in on_event:
|
|
edge_attrs['fontcolor'] = 'darkorange'
|
|
if 'fail' in on_event:
|
|
edge_attrs['fontcolor'] = 'red'
|
|
if 'success' in on_event:
|
|
edge_attrs['fontcolor'] = 'green'
|
|
else:
|
|
edge_attrs = {}
|
|
g.add_edge(pydot.Edge(nodes[start_state], nodes[end_state],
|
|
**edge_attrs))
|
|
|
|
start = pydot.Node("__start__", shape="point", width="0.1",
|
|
xlabel='start', fontcolor='green', **node_attrs)
|
|
g.add_node(start)
|
|
g.add_edge(pydot.Edge(start, nodes[source.default_start_state], style='dotted'))
|
|
|
|
print("*" * len(graph_name))
|
|
print(graph_name)
|
|
print("*" * len(graph_name))
|
|
print(source.pformat())
|
|
print(g.to_string().strip())
|
|
|
|
g.write(options.filename, format=options.format)
|
|
print("Created %s at '%s'" % (options.format, options.filename))
|
|
|
|
# To make the svg more pretty use the following:
|
|
# $ xsltproc ../diagram-tools/notugly.xsl ./states.svg > pretty-states.svg
|
|
# Get diagram-tools from https://github.com/vidarh/diagram-tools.git
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|