python-transitions-graph.py

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')