Merge pull request #2206 from pimotte/issue-2205

Making in_edges equivalent to out_edges

4 files changed, 65 insertions, 8 deletions

diff --git a/networkx/classes/digraph.py b/networkx/classes/digraph.py
index b209cbc6..f6c1a164 100644
--- a/networkx/classes/digraph.py
+++ b/networkx/classes/digraph.py
@@ -796,7 +796,7 @@ class DiGraph(Graph):
     # alias out_edges to edges
     out_edges = edges
 
-    def in_edges(self, nbunch=None, data=False):
+    def in_edges(self, nbunch=None, data=False, default=None):
         """Return an iterator over the incoming edges.
 
         Parameters
@@ -804,8 +804,15 @@ class DiGraph(Graph):
         nbunch : iterable container, optional (default= all nodes)
             A container of nodes.  The container will be iterated
             through once.
-        data : bool, optional (default=False)
-            If True, return edge attribute dict in 3-tuple (u,v,data).
+        data : string or bool, optional (default=False)
+            The edge attribute returned in 3-tuple (u,v,ddict[data]).
+            If True, return edge attribute dict in 3-tuple (u,v,ddict).
+            If False, return 2-tuple (u,v). 
+        default : value, optional (default=None)
+            Value used for edges that dont have the requested attribute.
+            Only relevant if data is not True or False.
+
+
 
         Returns
         -------
@@ -820,10 +827,15 @@ class DiGraph(Graph):
             nodes_nbrs=self.pred.items()
         else:
             nodes_nbrs=((n,self.pred[n]) for n in self.nbunch_iter(nbunch))
-        if data:
+        if data is True:
             for n,nbrs in nodes_nbrs:
                 for nbr,data in nbrs.items():
                     yield (nbr,n,data)
+        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 (nbr,n,d)
         else:
             for n,nbrs in nodes_nbrs:
                 for nbr in nbrs:
diff --git a/networkx/classes/multidigraph.py b/networkx/classes/multidigraph.py
index 2b850483..9eabdb87 100644
--- a/networkx/classes/multidigraph.py
+++ b/networkx/classes/multidigraph.py
@@ -493,7 +493,7 @@ class MultiDiGraph(MultiGraph,DiGraph):
     out_edges = edges
 
 
-    def in_edges(self, nbunch=None, data=False, keys=False):
+    def in_edges(self, nbunch=None, data=False, keys=False, default=None):
         """Return an iterator over the incoming edges.
 
         Parameters
@@ -501,10 +501,17 @@ class MultiDiGraph(MultiGraph,DiGraph):
         nbunch : iterable container, optional (default= all nodes)
             A container of nodes.  The container will be iterated
             through once.
-        data : bool, optional (default=False)
-            If True, return edge attribute dict with each edge.
+        data : string or bool, optional (default=False)
+            The edge attribute returned in 3-tuple (u,v,ddict[data]).
+            If True, return edge attribute dict in 3-tuple (u,v,ddict).
+            If False, return 2-tuple (u,v). 
+
         keys : bool, optional (default=False)
             If True, return edge keys with each edge.
+        default : value, optional (default=None)
+            Value used for edges that dont have the requested attribute.
+            Only relevant if data is not True or False.
+
 
         Returns
         -------
@@ -519,7 +526,7 @@ class MultiDiGraph(MultiGraph,DiGraph):
             nodes_nbrs = self.pred.items()
         else:
             nodes_nbrs = ((n, self.pred[n]) for n in self.nbunch_iter(nbunch))
-        if data:
+        if data is True:
             for n, nbrs in nodes_nbrs:
                 for nbr, keydict in nbrs.items():
                     for key, data in keydict.items():
@@ -527,6 +534,15 @@ class MultiDiGraph(MultiGraph,DiGraph):
                             yield (nbr, n, key, data)
                         else:
                             yield (nbr, n, data)
+        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
+                        if keys:
+                            yield (nbr, n, key, d)
+                        else:
+                            yield (nbr, n, d)
         else:
             for n, nbrs in nodes_nbrs:
                 for nbr, keydict in nbrs.items():
diff --git a/networkx/classes/tests/test_digraph.py b/networkx/classes/tests/test_digraph.py
index 339fa50c..4e725bd1 100644
--- a/networkx/classes/tests/test_digraph.py
+++ b/networkx/classes/tests/test_digraph.py
@@ -70,6 +70,13 @@ class BaseDiGraphTester(BaseGraphTester):
         assert_equal(sorted(G.out_edges(0)),[(0, 1)])
         assert_equal(sorted(G.out_edges(2)),[])
 
+    def test_out_edges_data(self):
+        G=networkx.DiGraph([(0, 1, {'data' : 0}), (1, 0, {})])
+        assert_equal(sorted(G.out_edges(data=True)), [(0, 1, {'data' : 0}), (1, 0, {})])
+        assert_equal(sorted(G.out_edges(0, data=True)), [(0, 1, {'data' : 0})])
+        assert_equal(sorted(G.out_edges(data='data')), [(0, 1, 0), (1, 0, None)])
+        assert_equal(sorted(G.out_edges(0, data='data')), [(0, 1, 0)])
+
     def test_in_edges_dir(self):
         G=self.P3
         assert_equal(sorted(G.in_edges()),[(0, 1), (1, 2)])
@@ -82,6 +89,13 @@ class BaseDiGraphTester(BaseGraphTester):
         assert_equal(sorted(G.in_edges(0)),[])
         assert_equal(sorted(G.in_edges(2)),[(1,2)])
 
+    def test_in_edges_data(self):
+        G=networkx.DiGraph([(0, 1, {'data' : 0}), (1, 0, {})])
+        assert_equal(sorted(G.in_edges(data=True)), [(0, 1, {'data' : 0}), (1, 0, {})])
+        assert_equal(sorted(G.in_edges(1, data=True)), [(0, 1, {'data' : 0})])
+        assert_equal(sorted(G.in_edges(data='data')), [(0, 1, 0), (1, 0, None)])
+        assert_equal(sorted(G.in_edges(1, data='data')), [(0, 1, 0)])
+
     def test_degree(self):
         G=self.K3
         assert_equal(list(G.degree()),[(0,4),(1,4),(2,4)])
diff --git a/networkx/classes/tests/test_multidigraph.py b/networkx/classes/tests/test_multidigraph.py
index a58e9532..72207da4 100644
--- a/networkx/classes/tests/test_multidigraph.py
+++ b/networkx/classes/tests/test_multidigraph.py
@@ -45,6 +45,14 @@ class BaseMultiDiGraphTester(BaseMultiGraphTester):
         assert_equal(sorted(G.out_edges()),
                      [(0,1),(0,1),(0,2),(1,0),(1,2),(2,0),(2,1)])
 
+    def test_out_edges_data(self):
+        G=self.K3
+        assert_equal(sorted(G.edges(0,data=True)),[(0,1,{}),(0,2,{})])
+        G.remove_edge(0,1)
+        G.add_edge(0,1,data=1)
+        assert_equal(sorted(G.edges(0,data=True)),[(0,1,{'data':1}),(0,2,{})])
+        assert_equal(sorted(G.edges(0,data='data')),[(0,1,1),(0,2,None)])
+
     def test_in_edges(self):
         G=self.K3
         assert_equal(sorted(G.in_edges()),
@@ -69,6 +77,13 @@ class BaseMultiDiGraphTester(BaseMultiGraphTester):
                      [(0,1,{}),(0,1,{}),(0,2,{}),(1,0,{}),(1,2,{}),
                       (2,0,{}),(2,1,{})])
 
+    def test_in_edges_data(self):
+        G=self.K3
+        assert_equal(sorted(G.in_edges(0,data=True)),[(1,0,{}),(2,0,{})])
+        G.remove_edge(1,0)
+        G.add_edge(1,0, data=1)
+        assert_equal(sorted(G.in_edges(0,data=True)),[(1,0,{'data':1}),(2,0,{})])
+        assert_equal(sorted(G.in_edges(0,data='data')),[(1,0,1),(2,0,None)])
 
     def is_shallow(self,H,G):
         # graph