#!/usr/bin/env python3 # 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. from unittest 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) from automaton.converters 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() self.selector = mock.MagicMock() def make_machine(start_state, transitions, event_name_cb): machine = machines.FiniteMachine() machine.add_state(start_state) machine.default_start_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) event = event_name_cb(start_state, end_state) machine.add_transition(start_state, end_state, event) return machine 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("--flow", dest="flow", action='store_true', help="use flow 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, options.flow, ] provided = sum([int(i) for i in types]) if provided > 1: parser.error("Only one of task/retry/engines/wbe requests/jobs/flow" " may be specified.") if provided == 0: parser.error("One of task/retry/engines/wbe requests/jobs/flow" " must be specified.") event_name_cb = lambda start_state, end_state: "on_%s" % end_state.lower() internal_states = list() ordering = 'in' if options.tasks: source_type = "Tasks" source = make_machine(states.PENDING, list(states._ALLOWED_TASK_TRANSITIONS), event_name_cb) elif options.retries: source_type = "Retries" source = make_machine(states.PENDING, list(states._ALLOWED_RETRY_TRANSITIONS), event_name_cb) elif options.flow: source_type = "Flow" source = make_machine(states.PENDING, list(states._ALLOWED_FLOW_TRANSITIONS), event_name_cb) 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), event_name_cb) elif options.jobs: source_type = "Jobs" source = make_machine(states.UNCLAIMED, list(states._ALLOWED_JOB_TRANSITIONS), event_name_cb) graph_attrs = { 'ordering': ordering, } graph_name = "%s states" % source_type def node_attrs_cb(state): node_color = None if state in internal_states: node_color = 'blue' if state in (states.FAILURE, states.REVERT_FAILURE): node_color = 'red' if state == states.REVERTED: node_color = 'darkorange' if state in (states.SUCCESS, states.COMPLETE): node_color = 'green' node_attrs = {} if node_color: node_attrs['fontcolor'] = node_color return node_attrs def edge_attrs_cb(start_state, on_event, end_state): edge_attrs = {} if options.engines: edge_attrs['label'] = on_event.replace("_", " ").strip() 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' return edge_attrs g = pydot.convert(source, graph_name, graph_attrs=graph_attrs, node_attrs_cb=node_attrs_cb, edge_attrs_cb=edge_attrs_cb) 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()