NodeView and EdgesView classes for NetworkX Graphs

graphviews.py

from __future__ import print_function

class NodeView(object):
    def __init__(self, graph, data=False):
        self.G = graph
        self.data = data
    def __iter__(self):
        if data:
            return iter(self.G.node.items())
        return iter(self.G.node)
    def __str__(self):
        return str(self.G.node.keys())
    def __repr__(self):
        return "NodeView({})".format(repr(self.G))
    def __contains__(self, node):
        return node in self.G.node
    def __len__(self):
        return len(self.G.node)

class EdgesView(object):
    # TODO: make a class for each type of edge-tuple
    def __init__(self, graph, nbunch=None, data=False, default=None):
        self.graph = graph
        self.nbunch = nbunch
        self.data = data
        self.default = default
        if nbunch is None:
            self.nodes_nbrs = graph.adj.items()
        else:
            self.nodes_nbrs = [(n, graph[n]) for n in graph.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):
        nodes = [n for n,nbrs in self.nodes_nbrs]
#         args = map(repr, (self.graph, nodes, self.data, self.default))
        return "EdgeView({}, nbunch={}, data={}, default={})".format(self.graph, nodes, self.data, self.default)
    def __contains__(self, u, v):
        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

if __name__ == "__main__":
    import networkx as nx
    G = nx.path_graph(9)
    print(EdgeView(G))
    print(NodeView(G))
    e=EdgeView(G, data=True)
    for i in e:
        print(i)
    
    G.edges = lambda nbunch, data, default: EdgeView(G, nbunch, data, default)
    print(G.edges((2,3), data=True))