directedgraph.py

class DirectedGraph:
    """A basic implementation of a directed graph."""

    def __init__(self):
        self.graph = {}

    def add_node(self, node):
        if node not in self.graph:
            self.graph[node] = {}

    def add_edge(self, start, end, payload={}):
        if start not in self.graph:
            self.add_node(start)
        if end not in self.graph:
            self.add_node(end)
        self.graph[start][end] = payload

    def nodes(self):
        return list(self.graph.keys())

    def edges(self):
        return list(self.graph.items())

    def parents(self, node):
        return [
            k for k, v in self.graph.items()
            if node in v.keys()
        ]

    def children(self, node):
        return list(self.graph[node])

    def descendants(self, start_node):
        visited = set()
        stack = [start_node]
        reachable = []

        while stack:
            node = stack.pop()
            if node not in visited:
                visited.add(node)
                reachable.append(node)
                for neighbor in list(self.graph[node]):
                    if neighbor not in visited:
                        stack.append(neighbor)

        return reachable

    def __getitem__(self, key):
        return self.graph[key]

    def __contains__(self, node):
        return node in self.graph