graph.py

class DirectedGraph:
    def __init__(self, edges={}):
        # copy/cast constructor
        if isinstance(edges, DirectedGraph): edges = edges.label
        self.label = {}
        self.vertex_in = {}
        self.vertex_out = {}
        for edge, label in edges.items():
            self[edge] = label

    def __repr__(self):
        return 'DirectedGraph({})'.format(dict(self.label))

    def __iter__(self):
        return iter(self.label)

    def items(self):
        return self.label.items()

    def __contains__(self, edge):
        if isinstance(edge, tuple) and len(edge) == 2:
            return edge in self.label
        else:
            raise Exception('Edge containment must be tested with tuples of 2 elements')

    def __getitem__(self, v):
        if isinstance(v, tuple):
            if len(v) != 2: raise Exception('Only binary edges are supported')
            return self.label[v]
        else:
            return set(self.vertex_out[v])

    def __setitem__(self, edge, label=1):
        if isinstance(edge, tuple):
            if len(edge) != 2: raise Exception('Only binary edges are supported')
            self.label[edge] = label
            self.vertex_out.setdefault(edge[0], set()).add(edge[1])
            self.vertex_in.setdefault(edge[1], set()).add(edge[0])
        else:
            raise Exception('Key must be a pair e1,e2')

    def __delitem__(self, edge):
        if isinstance(edge, tuple):
            if len(edge) != 2: raise Exception('Only binary edges are supported')
            del self.label[edge]
            self.vertex_out[edge[0]].remove(edge[1])
            if not self.vertex_out[edge[0]]: del self.vertex_out[edge[0]]
            self.vertex_in[edge[1]].remove(edge[0])
            if not self.vertex_in[edge[1]]: del self.vertex_in[edge[1]]
        else:
            # remove vertex:
            v1 = edge
            to_remove = set()
            if v1 in self.vertex_out:
                for v2 in self.vertex_out[v1]:
                    to_remove.add((v1, v2))
            if v1 in self.vertex_in:
                for v2 in self.vertex_in[v1]:
                    to_remove.add((v2, v1))
            for edge in to_remove:
                del self[edge]

class UndirectedGraph(DirectedGraph):
    def __init__(self, edges={}):
        super().__init__(edges)

    def __repr__(self):
        return 'UndirectedGraph({})'.format(dict(self.label))

    def __setitem__(self, edge, label=1):
        super().__setitem__(edge, label)
        super().__setitem__(tuple(reversed(edge)), label)

    def __delitem__(self, edge):
        super().__delitem__(edge)
        super().__delitem__(tuple(reversed(edge)))