Add new degree() implementation to graph.py

1 files changed, 102 insertions, 79 deletions

diff --git a/networkx/classes/graph.py b/networkx/classes/graph.py
index f0914e91..ac1fbea6 100644
--- a/networkx/classes/graph.py
+++ b/networkx/classes/graph.py
@@ -1224,92 +1224,115 @@ class Graph(object):
         return iter(self.adj.items())
 
     def degree(self, nbunch=None, weight=None):
-        """Return the degree of a node or nodes.
+        if nbunch in self:
+            nbrs = self.adj[nbunch]
+            if weight is None:
+                return len(nbrs) + (1 if nbunch in nbrs else 0) # handle self-loops
+            return sum(dd.get(weight, 1) for nbr,dd in nbrs.items()) +\
+                    (nbrs[nbunch].get(weight, 1) if nbunch in nbrs else 0)
 
-        The node degree is the number of edges adjacent to that node.
-
-        Parameters
-        ----------
-        nbunch : iterable container, optional (default=all nodes)
-            A container of nodes.  The container will be iterated
-            through once.
-
-        weight : string or None, optional (default=None)
-           The edge attribute that holds the numerical value used
-           as a weight.  If None, then each edge has weight 1.
-           The degree is the sum of the edge weights adjacent to the node.
-
-        Returns
-        -------
-        nd : dictionary, or number
-            A dictionary with nodes as keys and degree as values or
-            a number if a single node is specified.
-
-        Examples
-        --------
-        >>> G = nx.Graph()   # or DiGraph, MultiGraph, MultiDiGraph, etc
-        >>> G.add_path([0,1,2,3])
-        >>> G.degree(0)
-        1
-        >>> G.degree([0,1])
-        {0: 1, 1: 2}
-        >>> list(G.degree([0,1]).values())
-        [1, 2]
-
-        """
-        if nbunch in self:      # return a single node
-            return next(self.degree_iter(nbunch, weight))[1]
-        else:           # return a dict
-            return dict(self.degree_iter(nbunch, weight))
-
-    def degree_iter(self, nbunch=None, weight=None):
-        """Return an iterator for (node, degree).
-
-        The node degree is the number of edges adjacent to the node.
-
-        Parameters
-        ----------
-        nbunch : iterable container, optional (default=all nodes)
-            A container of nodes.  The container will be iterated
-            through once.
-
-        weight : string or None, optional (default=None)
-           The edge attribute that holds the numerical value used
-           as a weight.  If None, then each edge has weight 1.
-           The degree is the sum of the edge weights adjacent to the node.
-
-        Returns
-        -------
-        nd_iter : an iterator
-            The iterator returns two-tuples of (node, degree).
-
-        See Also
-        --------
-        degree
-
-        Examples
-        --------
-        >>> G = nx.Graph()   # or DiGraph, MultiGraph, MultiDiGraph, etc
-        >>> G.add_path([0,1,2,3])
-        >>> list(G.degree_iter(0)) # node 0 with degree 1
-        [(0, 1)]
-        >>> list(G.degree_iter([0,1]))
-        [(0, 1), (1, 2)]
-
-        """
         if nbunch is None:
             nodes_nbrs = self.adj.items()
         else:
             nodes_nbrs = ((n, self.adj[n]) for n in self.nbunch_iter(nbunch))
-
         if weight is None:
-            for n, nbrs in nodes_nbrs:
-                yield (n, len(nbrs) + (n in nbrs))  # return tuple (n,degree)
+            def d_iter():
+                for n, nbrs in nodes_nbrs:
+                    yield (n, len(nbrs) + (1 if n in nbrs else 0))  # return tuple (n,degree)
         else:
-            # edge weighted graph - degree is sum of nbr edge weights
-            for n, nbrs in nodes_nbrs:
-                yield (n, sum((nbrs[nbr].get(weight, 1) for nbr in nbrs)) +
-                    (n in nbrs and nbrs[n].get(weight, 1)))
+            def d_iter():
+                for n, nbrs in nodes_nbrs:
+                    yield (n, sum((nbrs[nbr].get(weight, 1) for nbr in nbrs)) +
+                        (nbrs[n].get(weight, 1) if n in nbrs else 0))
+        return d_iter()
+
+    # def degree(self, nbunch=None, weight=None):
+    #     """Return the degree of a node or nodes.
+
+    #     The node degree is the number of edges adjacent to that node.
+
+    #     Parameters
+    #     ----------
+    #     nbunch : iterable container, optional (default=all nodes)
+    #         A container of nodes.  The container will be iterated
+    #         through once.
+
+    #     weight : string or None, optional (default=None)
+    #        The edge attribute that holds the numerical value used
+    #        as a weight.  If None, then each edge has weight 1.
+    #        The degree is the sum of the edge weights adjacent to the node.
+
+    #     Returns
+    #     -------
+    #     nd : dictionary, or number
+    #         A dictionary with nodes as keys and degree as values or
+    #         a number if a single node is specified.
+
+    #     Examples
+    #     --------
+    #     >>> G = nx.Graph()   # or DiGraph, MultiGraph, MultiDiGraph, etc
+    #     >>> G.add_path([0,1,2,3])
+    #     >>> G.degree(0)
+    #     1
+    #     >>> G.degree([0,1])
+    #     {0: 1, 1: 2}
+    #     >>> list(G.degree([0,1]).values())
+    #     [1, 2]
+
+    #     """
+    #     if nbunch in self:      # return a single node
+    #         return next(self.degree_iter(nbunch, weight))[1]
+    #     else:           # return a dict
+    #         return dict(self.degree_iter(nbunch, weight))
+
+    # def degree_iter(self, nbunch=None, weight=None):
+    #     """Return an iterator for (node, degree).
+
+    #     The node degree is the number of edges adjacent to the node.
+
+    #     Parameters
+    #     ----------
+    #     nbunch : iterable container, optional (default=all nodes)
+    #         A container of nodes.  The container will be iterated
+    #         through once.
+
+    #     weight : string or None, optional (default=None)
+    #        The edge attribute that holds the numerical value used
+    #        as a weight.  If None, then each edge has weight 1.
+    #        The degree is the sum of the edge weights adjacent to the node.
+
+    #     Returns
+    #     -------
+    #     nd_iter : an iterator
+    #         The iterator returns two-tuples of (node, degree).
+
+    #     See Also
+    #     --------
+    #     degree
+
+    #     Examples
+    #     --------
+    #     >>> G = nx.Graph()   # or DiGraph, MultiGraph, MultiDiGraph, etc
+    #     >>> G.add_path([0,1,2,3])
+    #     >>> list(G.degree_iter(0)) # node 0 with degree 1
+    #     [(0, 1)]
+    #     >>> list(G.degree_iter([0,1]))
+    #     [(0, 1), (1, 2)]
+
+    #     """
+    #     if nbunch is None:
+    #         nodes_nbrs = self.adj.items()
+    #     else:
+    #         nodes_nbrs = ((n, self.adj[n]) for n in self.nbunch_iter(nbunch))
+
+    #     if weight is None:
+    #         for n, nbrs in nodes_nbrs:
+    #             yield (n, len(nbrs) + (n in nbrs))  # return tuple (n,degree)
+    #     else:
+    #         # edge weighted graph - degree is sum of nbr edge weights
+    #         for n, nbrs in nodes_nbrs:
+    #             yield (n, sum((nbrs[nbr].get(weight, 1) for nbr in nbrs)) +
+    #                 (n in nbrs and nbrs[n].get(weight, 1)))
 
     def clear(self):
         """Remove all nodes and edges from the graph.