ENH: let users set a default value in get_attr methods (networkx#6887)

* ENH: let users set a default value in get_attr methods

Co-authored-by: Dan Schult <dschult@colgate.edu>

doc/release/release_dev.rst

@@ -19,7 +19,6 @@ Highlights
 This release is the result of X of work with over X pull requests by
 X contributors. Highlights include:
 
-
 Improvements
 ------------
 
@@ -42,6 +41,10 @@ API Changes
   In `is_semiconnected`, the keyword argument `topo_order` has been removed.
   That argument resulted in silently incorrect results more often than not.
 
+- [`#6887 <https://github.com/networkx/networkx/pull/6887>`_]
+  A new ``default`` argument is added to ``get_node_attributes`` and
+  ``get_edge_attributes``. The ``default`` keyword can be used to set
+  a default value if the attribute is missing from a node/edge.
 
 
 Deprecations

networkx/classes/function.py

@@ -649,7 +649,7 @@ def set_node_attributes(G, values, name=None):
                 pass
 
 
-def get_node_attributes(G, name):
+def get_node_attributes(G, name, default=None):
     """Get node attributes from graph
 
     Parameters
@@ -659,6 +659,11 @@ def get_node_attributes(G, name):
     name : string
        Attribute name
 
+    default: object (default=None)
+       Default value of the node attribute if there is no value set for that
+       node in graph. If `None` then nodes without this attribute are not
+       included in the returned dict.
+
     Returns
     -------
     Dictionary of attributes keyed by node.
@@ -670,7 +675,13 @@ def get_node_attributes(G, name):
     >>> color = nx.get_node_attributes(G, "color")
     >>> color[1]
     'red'
+    >>> G.add_node(4)
+    >>> color = nx.get_node_attributes(G, "color", default="yellow")
+    >>> color[4]
+    'yellow'
     """
+    if default is not None:
+        return {n: d.get(name, default) for n, d in G.nodes.items()}
     return {n: d[name] for n, d in G.nodes.items() if name in d}
 
 
@@ -811,7 +822,7 @@ def set_edge_attributes(G, values, name=None):
                     pass
 
 
-def get_edge_attributes(G, name):
+def get_edge_attributes(G, name, default=None):
     """Get edge attributes from graph
 
     Parameters
@@ -821,6 +832,11 @@ def get_edge_attributes(G, name):
     name : string
        Attribute name
 
+    default: object (default=None)
+       Default value of the edge attribute if there is no value set for that
+       edge in graph. If `None` then edges without this attribute are not
+       included in the returned dict.
+
     Returns
     -------
     Dictionary of attributes keyed by edge. For (di)graphs, the keys are
@@ -834,11 +850,17 @@ def get_edge_attributes(G, name):
     >>> color = nx.get_edge_attributes(G, "color")
     >>> color[(1, 2)]
     'red'
+    >>> G.add_edge(3, 4)
+    >>> color = nx.get_edge_attributes(G, "color", default="yellow")
+    >>> color[(3, 4)]
+    'yellow'
     """
     if G.is_multigraph():
         edges = G.edges(keys=True, data=True)
     else:
         edges = G.edges(data=True)
+    if default is not None:
+        return {x[:-1]: x[-1].get(name, default) for x in edges}
     return {x[:-1]: x[-1][name] for x in edges if name in x[-1]}
 
 

networkx/classes/tests/test_function.py

@@ -623,6 +623,10 @@ def test_get_node_attributes():
         assert attrs[0] == vals
         assert attrs[1] == vals
         assert attrs[2] == vals
+        default_val = 1
+        G.add_node(4)
+        attrs = nx.get_node_attributes(G, attr, default=default_val)
+        assert attrs[4] == default_val
 
 
 def test_get_edge_attributes():
@@ -633,22 +637,18 @@ def test_get_edge_attributes():
         vals = 100
         nx.set_edge_attributes(G, vals, attr)
         attrs = nx.get_edge_attributes(G, attr)
-
         assert len(attrs) == 2
-        if G.is_multigraph():
-            keys = [(0, 1, 0), (1, 2, 0)]
-            for u, v, k in keys:
-                try:
-                    assert attrs[(u, v, k)] == 100
-                except KeyError:
-                    assert attrs[(v, u, k)] == 100
-        else:
-            keys = [(0, 1), (1, 2)]
-            for u, v in keys:
-                try:
-                    assert attrs[(u, v)] == 100
-                except KeyError:
-                    assert attrs[(v, u)] == 100
+
+        for edge in G.edges:
+            assert attrs[edge] == vals
+
+        default_val = vals
+        G.add_edge(4, 5)
+        deafult_attrs = nx.get_edge_attributes(G, attr, default=default_val)
+        assert len(deafult_attrs) == 3
+
+        for edge in G.edges:
+            assert deafult_attrs[edge] == vals
 
 
 def test_is_empty():