Finding all successors of a node with pygraphviz

pygraphviz.AGraph class implements the successors method for finding first level successors of a given node.

g.successors('a') returns ['b', 'c', 'd', 'z']

Sometimes a need arises to find all the successors of all successive nodes as well. The class can be extended to allow for this behavior. The following example illustrates this by extending the base class with all_successors method

#!/usr/bin/python

import pygraphviz as pgv

class Graph(pgv.AGraph):
    def __init__(self, **kvargs):
        pgv.AGraph.__init__(self, **kvargs)
        self.__successors = None
    def all_successors(self, node):
        self.__successors = []
        start_node = self.get_node(node)
        def recur_successor(node):
            node_successors = self.successors(node)
            for s in node_successors:
                if s not in self.__successors:
                    self.__successors.append(s)
                    recur_successor(s)

        recur_successor(start_node)

        return self.__successors

g = Graph(directed=True)

g.add_edge('a', 'b')
g.add_edge('a', 'c')
g.add_edge('a', 'd')
g.add_edge('d', 'e')
g.add_edge('e', 'f')
g.add_edge('x', 'b')
g.add_edge('a', 'z')
g.add_edge('r', 's')
g.add_edge('b','c')

g.draw('s1.png', prog='fdp')

print g.successors('a')
print g.all_successors('a')

related_to_a = g.all_successors('a')
related_to_a.extend('a')
for n in g.nodes():
    if n not in related_to_a:
        g.delete_node(n)

g.draw('s2.png', prog='fdp')




g.all_successors('a') returns ['b', 'c', 'd', 'e', 'f', 'z']

'r' and 's' are not related to 'a' at all and 'x', although it's a predecessor of 'b' it is not a successor of 'a'.