summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorMridul Seth <seth.mridul@gmail.com>2015-12-28 18:41:32 +0530
committerMridul Seth <seth.mridul@gmail.com>2015-12-28 18:41:32 +0530
commitdb740e9e699175fb6510a3bdad0c6b2f63e90f11 (patch)
tree8c91da29511cef83dabdcf7ac05fe302fc7fe629
parentbe5b6adde1bb27b8ef01687af854754f344868b2 (diff)
downloadnetworkx-db740e9e699175fb6510a3bdad0c6b2f63e90f11.tar.gz
add option to view specific attributes in G.nodes()
-rw-r--r--networkx/classes/graph.py46
-rw-r--r--networkx/classes/tests/test_graph.py15
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