NodeView and EdgesView classes for NetworkX Graphs

graphview.py

from __future__ import print_function

class NodesView(object):
    def __init__(self, graph, data=False, default=None):
        self.G = graph
        self.data = data
        self.default = default
    def __iter__(self):
        data=self.data
        if data is True:
            return iter(self.G.node.items())
        if data is False:
            return iter(self.G.node)
        return ( (n,dd.get(data, self.default)) for n,dd in self.G.node.items() )
    def __str__(self):
        if self.data is False:
            return str(self.G.node.keys())
        if self.data is True:
            return str(self.G.node.items())
        return str(dict(self))
    def __repr__(self):
        args = (repr(self.G), self.data, self.default)
        return "NodesView({}, data={}, default={})".format(*args)
    def __contains__(self, n):
        return n in self.G.node
    def __len__(self):
        return len(self.G.node)
    def __and__(self, other):
        return set(self.G.node) & set(other)
    def __or__(self, other):
        return set(self.G.node) | set(other)
    def __xor__(self, other):
        return set(self.G.node) ^ set(other)
    def __sub__(self, other):
        return set(self.G.node) - set(other)


class EdgesView(object):
    # TODO: make a class for each type of edge-tuple
    def __init__(self, graph, nbunch=None, data=False, default=None):
        self.G = G = graph
        self.nbunch = nbunch
        self.data = data
        self.default = default
        if nbunch is None:
            self.nodes_nbrs = G.adj.items()
        else:
            self.nodes_nbrs = [(n,G[n]) for n in G.nbunch_iter(nbunch)]
    def __iter__(self):
        seen = {}
        data = self.data
        default = self.default
        nodes_nbrs = self.nodes_nbrs
        if data is True:
            for n, nbrs in nodes_nbrs:
                for nbr, ddict in nbrs.items():
                    if nbr not in seen:
                        yield (n, nbr, ddict)
                seen[n] = 1
        elif data is not False:
            for n, nbrs in nodes_nbrs:
                for nbr, ddict in nbrs.items():
                    if nbr not in seen:
                        d = ddict[data] if data in ddict else default
                        yield (n, nbr, d)
                seen[n] = 1
        else:  # data is False
            for n, nbrs in nodes_nbrs:
                for nbr in nbrs:
                    if nbr not in seen:
                        yield (n, nbr)
                seen[n] = 1
        del seen
    def __str__(self):
        return str(list(self))
    def __repr__(self):
        args = map(repr, (self.G, self.nbunch, self.data, self.default))
        msg="EdgesView({}, nbunch={}, data={}, default={})"
        return msg.format(*args)
    def __contains__(self, e):
        u,v = e
        try:
            return v in self.G.adj[u]
        except KeyError:
            return False
    def __len__(self):
        return sum(len(nbrs) for n, nbrs in self.nodes_nbrs)//2
    def __and__(self, other):
        return set(self) & set(other)
    def __or__(self, other):
        return set(self) | set(other)
    def __xor__(self, other):
        return set(self) ^ set(other)
    def __sub__(self, other):
        return set(self) - set(other)

UndirectedEdgesView = EdgesView

class DirectedEdgesView(EdgesView):
    def __iter__(self):
        data = self.data
        default = self.default
        nodes_nbrs = self.nodes_nbrs
        if data is True:
            for n, nbrs in nodes_nbrs:
                for nbr, ddict in nbrs.items():
                    yield (n, nbr, ddict)
        elif data is not False:
            for n, nbrs in nodes_nbrs:
                for nbr, ddict in nbrs.items():
                    d = ddict[data] if data in ddict else default
                    yield (n, nbr, d)
        else:  # data is False
            for n, nbrs in nodes_nbrs:
                for nbr in nbrs:
                    yield (n, nbr)
    def __repr__(self):
        args = map(repr, (self.G, self.nbunch, self.data, self.default))
        msg="DirectedEdgesView({}, nbunch={}, data={}, default={})"
        return msg.format(*args)

class MultiEdgesView(EdgesView):
    def __init__(self, graph, keys=False, **kwds):
        super(self).__init__(graph, **kwds)
        self.keys = keys
    def __iter__(self):
        seen = {}
        data = self.data
        default = self.default
        nodes_nbrs = self.nodes_nbrs
        keys = self.keys
        if data is True:
            for n, nbrs in nodes_nbrs:
                for nbr, keydict in nbrs.items():
                    if nbr not in seen:
                        for key, ddict in keydict.items():
                            yield (n, nbr, key, ddict) if keys else (n, nbr, ddict)
                seen[n] = 1
        elif data is not False:
            for n, nbrs in nodes_nbrs:
                for nbr, keydict in nbrs.items():
                    if nbr not in seen:
                        for key, ddict in keydict.items():
                            d = ddict[data] if data in ddict else default
                            yield (n, nbr, key, d) if keys else (n, nbr, d)
                seen[n] = 1
        else:
            for n, nbrs in nodes_nbrs:
                for nbr, keydict in nbrs.items():
                    if nbr not in seen:
                        for key in keydict:
                            yield (n, nbr, key) if keys else (n, nbr)
                seen[n] = 1
        del seen
    def __repr__(self):
        args = (self.G, self.nbunch, self.data, self.default, self.keys)
        args = map(repr, args)
        msg = ("MultiEdgesView({}, nbunch={}, data={}, "
              "default={}, keys={})")
        return msg.format(*args)

UndirectedMultiEdgesView = MultiEdgesView

class DirectedMultiEdgesView(MultiEdgesView):
    def __iter__(self):
        data = self.data
        default = self.default
        nodes_nbrs = self.nodes_nbrs
        keys = self.keys
        if data is True:
            for n, nbrs in nodes_nbrs:
                for nbr, keydict in nbrs.items():
                    for key, ddict in keydict.items():
                        yield (n, nbr, key, ddict) if keys else (n, nbr, ddict)
        elif data is not False:
            for n, nbrs in nodes_nbrs:
                for nbr, keydict in nbrs.items():
                    for key, ddict in keydict.items():
                        d = ddict[data] if data in ddict else default
                        yield (n, nbr, key, d) if keys else (n, nbr, d)
        else:
            for n, nbrs in nodes_nbrs:
                for nbr, keydict in nbrs.items():
                    for key in keydict:
                        yield (n, nbr, key) if keys else (n, nbr)
    def __repr__(self):
        args = (self.G, self.nbunch, self.data, self.default, self.keys)
        args = map(repr, args)
        msg = ("DirectedMultiEdgesView({}, nbunch={}, data={}, "
                                             "default={}, keys={})")
        return msg.format(*args)

if __name__ == "__main__":
    import networkx as nx
    G = nx.path_graph(9)
    print("Nodes:", NodesView(G))
    print("Edges:", EdgesView(G))
    print("repr Nodes:", repr(NodesView(G)))
    print("repr Edges:", repr(EdgesView(G)))
    print("Edges with data:")
    for e in EdgesView(G, data=True):
        print(e)

    # hack an example of this in place of G.edges
    G.edges = lambda nbunch=None, data=False, default=None: \
        EdgesView(G, nbunch, data, default)

    print("Print edges from nodes 2 and 3 with data")
    print(G.edges((2,3), data=True))
    print()
    H = nx.Graph()
    H.add_path(range(7,14))
    gnv = EdgesView(G)
    hnv = EdgesView(H)
    print("G edges:", gnv)
    print("H edges:", hnv)
    print("G & H edges:", gnv & hnv)
    print("G - H edges:", gnv - hnv)
    print("G | H edges:", gnv | hnv)
    print("G ^ H edges:", gnv ^ hnv)