diff options
author | Mridul Seth <seth.mridul@gmail.com> | 2015-12-28 18:41:32 +0530 |
---|---|---|
committer | Mridul Seth <seth.mridul@gmail.com> | 2015-12-28 18:41:32 +0530 |
commit | db740e9e699175fb6510a3bdad0c6b2f63e90f11 (patch) | |
tree | 8c91da29511cef83dabdcf7ac05fe302fc7fe629 | |
parent | be5b6adde1bb27b8ef01687af854754f344868b2 (diff) | |
download | networkx-db740e9e699175fb6510a3bdad0c6b2f63e90f11.tar.gz |
add option to view specific attributes in G.nodes()
-rw-r--r-- | networkx/classes/graph.py | 46 | ||||
-rw-r--r-- | networkx/classes/tests/test_graph.py | 15 |
2 files changed, 45 insertions, 16 deletions
diff --git a/networkx/classes/graph.py b/networkx/classes/graph.py index b12b960a..6d785933 100644 --- a/networkx/classes/graph.py +++ b/networkx/classes/graph.py @@ -608,22 +608,30 @@ class Graph(object): except KeyError: pass - def nodes(self, data=False): + def nodes(self, data=False, default=None): """Returns an iterator over the nodes. Parameters ---------- data : boolean, optional (default=False) - If ``False``, the iterator returns nodes. If ``True``, - the iterator return a two-tuple of node and node data - dictionary. + If ``False``, the iterator returns nodes. If ``True``, + the iterator returns a two-tuple of node and node data + dictionary. ``data`` can also be a specific node attribute from + node data dictionary, if ``data`` is an attribute, the iterator + returns a two-tuple of node and value of node attribute. + + default : value, optional (default=None) + Value used for nodes that dont have the requested attribute. + Only relevant if data is not True or False. Returns ------- iterator An iterator over nodes, or if ``data`` is ``True``, an iterator over two-tuples of the form ``(node, node data - dictionary)``. + dictionary)``, or if ``data`` is an attribute from the node + data dictionary, an iterator over two-tuples of the form ``(node, + attribute value)`` Notes ----- @@ -632,7 +640,8 @@ class Graph(object): Examples -------- - There are two simple ways of getting a list of all nodes in the graph:: + There are two simple ways of getting a list of all nodes in the graph: + >>> G = nx.Graph() >>> G.add_nodes_from(range(3)) >>> list(G.nodes()) @@ -640,16 +649,31 @@ class Graph(object): >>> list(G) [0, 1, 2] - To get the node data along with the nodes:: + To get the node data along with the nodes: + >>> G.add_node(1, time='5pm') >>> G.node[0]['foo'] = 'bar' >>> list(G.nodes(data=True)) [(0, {'foo': 'bar'}), (1, {'time': '5pm'}), (2, {})] - + >>> list(G.nodes(data='foo')) + [(0, 'bar'), (1, None), (2, None)] + >>> list(G.nodes(data='time')) + [(0, None), (1, '5pm'), (2, None)] + >>> list(G.nodes(data='time', default='Not Available')) + [(0, 'Not Available'), (1, '5pm'), (2, 'Not Available')] """ - if data: - return iter(self.node.items()) - return iter(self.node) + ndict = self.node.items() + if data is True: + for n, ddict in ndict: + yield (n, ddict) + elif data is not False: + for n, ddict in ndict: + d = ddict[data] if data in ddict else default + yield (n, d) + else: + for n in self.node: + yield n + def number_of_nodes(self): """Return the number of nodes in the graph. diff --git a/networkx/classes/tests/test_graph.py b/networkx/classes/tests/test_graph.py index 982c96ea..c54b4965 100644 --- a/networkx/classes/tests/test_graph.py +++ b/networkx/classes/tests/test_graph.py @@ -1,5 +1,5 @@ -#!/usr/bin/env python -from nose.tools import * +from nose.tools import (assert_equal, assert_raises, assert_true, raises, + assert_not_equal) import networkx class BaseGraphTester(object): @@ -324,14 +324,19 @@ class BaseAttrGraphTester(BaseGraphTester): assert_equal(H.graph['foo'], 'bar') def test_node_attr(self): - G=self.K3 - G.add_node(1,foo='bar') + G = self.K3 + G.add_node(1, foo='bar') assert_equal(list(G.nodes()), [0, 1, 2]) assert_equal(list(G.nodes(data=True)), [(0, {}), (1, {'foo':'bar'}), (2, {})]) - G.node[1]['foo']='baz' + G.node[1]['foo'] = 'baz' assert_equal(list(G.nodes(data=True)), [(0, {}), (1, {'foo':'baz'}), (2, {})]) + assert_equal(list(G.nodes(data='foo')), + [(0, None), (1, 'baz'), (2, None)]) + assert_equal(list(G.nodes(data='foo', default='bar')), + [(0, 'bar'), (1, 'baz'), (2, 'bar')]) + def test_node_attr2(self): G=self.K3 |