updating viz module and all tests to networkx 2.0

brooksandrew · Sep 24, 2017 · e8a886b · e8a886b
1 parent 8af8818
commit e8a886b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/postman_problems/tests/test_example_sleeping_giant.py b/postman_problems/tests/test_example_sleeping_giant.py
@@ -58,10 +58,10 @@ def test_add_node_attributes():
         assert 'Y' in v
 
     # spot check node attributes for first node
-    print(graph_node_attrs.nodes(data=True)[0][0])
-    assert graph_node_attrs.nodes(data=True)[0][0] == 'rs_end_north'
-    assert graph_node_attrs.nodes(data=True)[0][1]['X'] == 1772
-    assert graph_node_attrs.nodes(data=True)[0][1]['Y'] == 172
+    node_data_from_graph = list(graph_node_attrs.nodes(data=True))
+    assert node_data_from_graph[0][0] == 'rs_end_north'
+    assert node_data_from_graph[0][1]['X'] == 1772
+    assert node_data_from_graph[0][1]['Y'] == 172
 
 
 def test_get_shortest_paths_distances():

diff --git a/postman_problems/viz.py b/postman_problems/viz.py
@@ -3,6 +3,7 @@
 import imageio
 import tqdm
 import numpy as np
+import networkx as nx
 import graphviz as gv
 from collections import defaultdict
 
@@ -22,7 +23,7 @@ def add_node_attributes(graph, nodelist):
         networkx graph: original `graph` augmented w node attributes
     """
     for i, row in nodelist.iterrows():
-        graph.node[row['id']] = row.to_dict()
+        nx.set_node_attributes(graph, {row['id']: row.to_dict()})
     return graph