view python/graph.py @ 284:05184b95fa16

Moved tests to seperate folder
author Windel Bouwman
date Fri, 15 Nov 2013 13:43:22 +0100
parents 9fca39eebe50
children
line wrap: on
line source


class Graph:
    """
       Generic graph base class.
       Can dump to graphiz dot format for example!
    """
    def __init__(self):
        self.nodes = set()
        self.edges = set()

    def addNode(self, n):
        self.nodes.add(n)

    def delNode(self, n):
        self.nodes.remove(n)

    def addEdge(self, n, m):
        """ Add an edge from n to m """
        self.edges.add((n, m))
        self.edges.add((m, n))

    def hasEdge(self, n, m):
        return (n, m) in self.edges

    def delEdge(self, n, m):
        self.edges.remove((n, m))
        self.edges.remove((m, n))

    def adjecent(self, n):
        """ Return all nodes with edges to n """
        # TODO: this can be optimized for speed:
        return set(m for m in self.nodes if self.hasEdge(n, m))
        
    def to_dot(self, f):
        """ Generate graphviz dot representation """
        for node in self.nodes:
            print('{} [label="{}" shape=box3d];'.format(id(node), node), file=f)
        for n, m in self.edges:
            print('{} -> {};'.format(id(n), id(m)), file=f)


class Node:
    """
       Node in a graph.
    """
    def __init__(self, g):
        self.g = g
        self.addDegree = 0 # Hack to increase degree

    @property
    def Adjecent(self):
        return self.g.adjecent(self)

    @property
    def Degree(self):
        return len(self.Adjecent) + self.addDegree


class DiGraph(Graph):
    """
        Directed graph.
    """
    def addEdge(self, n, m):
        """ Add an edge from n to m """
        self.edges.add((n, m))

    def hasEdge(self, n, m):
        return (n, m) in self.edges

    def successors(self, n):
        return set(m for m in self.nodes if self.hasEdge(n, m))

    def predecessors(self, n):
        return set(m for m in self.nodes if self.hasEdge(m, n))


class DiNode(Node):
    @property
    def Succ(self):
        return self.g.successors(self)

    @property
    def Pred(self):
        return self.g.predecessors(self)