openstack/automaton
Code Issues Proposed changes

Merge "Allow providing a state-name translation dict"

Browse Source
This commit is contained in:
Jenkins 2015-07-30 20:51:49 +00:00 committed by Gerrit Code Review
commit 183882193c
1 changed files with 11 additions and 3 deletions

14
automaton/converters/pydot.py

@ -25,7 +25,7 @@ except ImportError:
def convert(machine, graph_name, def convert(machine, graph_name,
graph_attrs=None, node_attrs_cb=None, edge_attrs_cb=None, graph_attrs=None, node_attrs_cb=None, edge_attrs_cb=None,
add_start_state=True): add_start_state=True, name_translations=None):
"""Translates the state machine into a pydot graph. """Translates the state machine into a pydot graph.
:param machine: state machine to convert :param machine: state machine to convert
@ -55,11 +55,16 @@ def convert(machine, graph_name,
``default_start_state`` then this does nothing, ``default_start_state`` then this does nothing,
even if enabled) even if enabled)
:type add_start_state: bool :type add_start_state: bool
:param name_translations: a dict that provides alternative ``state``
string names for each state
:type name_translations: dict
""" """
if not PYDOT_AVAILABLE: if not PYDOT_AVAILABLE:
raise RuntimeError("pydot (or pydot2 or equivalent) is required" raise RuntimeError("pydot (or pydot2 or equivalent) is required"
" to convert a state machine into a pydot" " to convert a state machine into a pydot"
" graph") " graph")
if not name_translations:
name_translations = {}
graph_kwargs = { graph_kwargs = {
'rankdir': 'LR', 'rankdir': 'LR',
'nodesep': '0.25', 'nodesep': '0.25',
@ -82,14 +87,17 @@ def convert(machine, graph_name,
start_node_attrs = node_attrs.copy() start_node_attrs = node_attrs.copy()
if node_attrs_cb is not None: if node_attrs_cb is not None:
start_node_attrs.update(node_attrs_cb(start_state)) start_node_attrs.update(node_attrs_cb(start_state))
nodes[start_state] = pydot.Node(start_state, pretty_start_state = name_translations.get(start_state,
start_state)
nodes[start_state] = pydot.Node(pretty_start_state,
**start_node_attrs) **start_node_attrs)
g.add_node(nodes[start_state]) g.add_node(nodes[start_state])
if end_state not in nodes: if end_state not in nodes:
end_node_attrs = node_attrs.copy() end_node_attrs = node_attrs.copy()
if node_attrs_cb is not None: if node_attrs_cb is not None:
end_node_attrs.update(node_attrs_cb(end_state)) end_node_attrs.update(node_attrs_cb(end_state))
nodes[end_state] = pydot.Node(end_state, **end_node_attrs) pretty_end_state = name_translations.get(end_state, end_state)
nodes[end_state] = pydot.Node(pretty_end_state, **end_node_attrs)
g.add_node(nodes[end_state]) g.add_node(nodes[end_state])
edge_attrs = {} edge_attrs = {}
if edge_attrs_cb is not None: if edge_attrs_cb is not None: