Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import pygraphviz as pgv
class MachineGraph(object):
def get_graph(self, title=None):
"""Generate a DOT graph with pygraphviz."""
state_attrs = {
'shape': 'circle',
'height': '1.2',
}
machine_attrs = {
'directed': True,
'strict': False,
'rankdir': 'LR',
'ratio': '0.3'
}
fsm = self.machine
if title is None:
title = self.__class__.__name__
elif title is False:
title = ''
fsm_graph = pgv.AGraph(title=title, **machine_attrs)
fsm_graph.node_attr.update(state_attrs)
# for each state, draw a circle
for state in fsm.states.items():
shape = state_attrs['shape']
# we want the first state to be a double circle (UML style)
if state == fsm.states.items()[0]:
shape = 'doublecircle'
else:
shape = state_attrs['shape']
state = state[0]
fsm_graph.add_node(n=state, shape=shape)
fsm_graph.add_node('null', shape='plaintext', label='')
fsm_graph.add_edge('null', 'new')
# For each event, add the transitions
for event in fsm.events.items():
event = event[1]
label = str(event.name)
for transition in event.transitions.items():
src = transition[0]
dst = transition[1][0].dest
fsm_graph.add_edge(src, dst, label=label)
return fsm_graph
# extend your class with this class, create an instance and run:
# inst.get_graph().draw('graphs/foobar.png', prog='dot')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment