Changes nsl.tools.add_edge() to return a boolean result.

PiperOrigin-RevId: 317868388

Author: Neural-Link Team

neural_structured_learning/tools/graph_utils.py

@@ -63,15 +63,16 @@ def add_edge(graph, edge):
       supplied, it defaults to 1.0.
 
   Returns:
-    `None`. Instead, this function has a side-effect on the `graph` argument.
+    `True` if and only if a new edge was added to `graph`.
   """
   source = edge[0]
-  if source not in graph: graph[source] = {}
-  t_dict = graph[source]
   target = edge[1]
   weight = float(edge[2]) if len(edge) > 2 else 1.0
-  if target not in t_dict or weight > t_dict[target]:
+  t_dict = graph.setdefault(source, {})
+  is_new_edge = target not in t_dict
+  if is_new_edge or weight > t_dict[target]:
     t_dict[target] = weight
+  return is_new_edge
 
 
 def add_undirected_edges(graph):

neural_structured_learning/tools/graph_utils_test.py

@@ -29,13 +29,15 @@ class GraphUtilsTest(absltest.TestCase):
 
   def testAddEdge(self):
     graph = {}
-    graph_utils.add_edge(graph, ['A', 'B', '0.5'])
-    graph_utils.add_edge(graph, ['A', 'C', 0.7])  # Tests that the edge
-    graph_utils.add_edge(graph, ['A', 'C', 0.9])  # ...with maximal weight
-    graph_utils.add_edge(graph, ['A', 'C', 0.8])  # ...is used.
-    graph_utils.add_edge(graph, ('B', 'A', '0.4'))
-    graph_utils.add_edge(graph, ('B', 'C'))  # Tests default weight
-    graph_utils.add_edge(graph, ('D', 'A', 0.75))
+    self.assertTrue(graph_utils.add_edge(graph, ['A', 'B', '0.5']))
+    # The next 3 calls test that the edge with maximal weight is used.
+    self.assertTrue(graph_utils.add_edge(graph, ['A', 'C', 0.7]))
+    self.assertFalse(graph_utils.add_edge(graph, ['A', 'C', 0.9]))
+    self.assertFalse(graph_utils.add_edge(graph, ['A', 'C', 0.8]))
+    self.assertTrue(graph_utils.add_edge(graph, ('B', 'A', '0.4')))
+    # Tests that when no weight is specified, it defaults to 1.0.
+    self.assertTrue(graph_utils.add_edge(graph, ('B', 'C')))
+    self.assertTrue(graph_utils.add_edge(graph, ('D', 'A', 0.75)))
     self.assertDictEqual(graph, GRAPH)
 
   def testAddUndirectedEdges(self):