Merge pull request #35 from pygraphviz/attr-update

Attr update

pygraphviz/agraph.py

@@ -293,7 +293,7 @@ def add_node(self, n, **attr):
         >>> G.nodes()  # doctest: +IGNORE_UNICODE
         [u'a', u'1']
 
-        Attributes can be added to nodes on creation
+        Attributes can be added to nodes on creation or updated after creation
         (attribute values must be strings)
 
         >>> G.add_node(2,color='red')
@@ -302,7 +302,6 @@ def add_node(self, n, **attr):
         for a list of attributes.
 
         Anonymous Graphviz nodes are currently not implemented.
-
         """
         if not is_string_like(n):
             n = str(n)
@@ -311,8 +310,8 @@ def add_node(self, n, **attr):
             nh = gv.agnode(self.handle, n, _Action.find)
         except KeyError:
             nh = gv.agnode(self.handle, n, _Action.create)
-            node = Node(self, nh=nh)
-            node.attr.update(**attr)
+        node = Node(self, nh=nh)
+        node.attr.update(**attr)
 
     def add_nodes_from(self, nbunch, **attr):
         """Add nodes from a container nbunch.
@@ -326,10 +325,9 @@ def add_nodes_from(self, nbunch, **attr):
         [u'1', u'a', u'b', u'spam']
 
 
-        Attributes can be added to nodes on creation
+        Attributes can be added to nodes on creation or updated after creation
 
         >>> G.add_nodes_from(nlist, color='red') # set all nodes in nlist red
-
         """
         for n in nbunch:
             self.add_node(n, **attr)
@@ -343,7 +341,6 @@ def remove_node(self, n):
         >>> G=AGraph()
         >>> G.add_node('a')
         >>> G.remove_node('a')
-
         """
         if not is_string_like(n):
             n = str(n)
@@ -450,7 +447,7 @@ def add_edge(self, u, v=None, key=None, **attr):
         >>> sorted(G.edges(keys=True))  # doctest: +IGNORE_UNICODE
         [(u'a', u'b', u'first'), (u'a', u'b', u'second')]
 
-        Attributes can be added when edges are created
+        Attributes can be added when edges are created or updated after creation
 
         >>> G.add_edge('a','b',color='green')
 
@@ -472,16 +469,18 @@ def add_edge(self, u, v=None, key=None, **attr):
         except:
             self.add_node(v)
             vh = Node(self, v).handle
+        if key is not None:
+            if not is_string_like(key):
+                key = str(key)
+            key = key.encode(self.encoding)
         try:
-            if key is not None:
-                if not is_string_like(key):
-                    key = str(key)
-                key = key.encode(self.encoding)
+            # new
             eh = gv.agedge(self.handle, uh, vh, key, _Action.create)
-            e = Edge(self, eh=eh)
-            e.attr.update(**attr)
         except KeyError:
-            return None # silent failure for strict graph, already added
+            # for strict graph, or already added
+            eh = gv.agedge(self.handle, uh, vh, key, _Action.find)
+        e = Edge(self, eh=eh)
+        e.attr.update(**attr)
 
     def add_edges_from(self, ebunch, **attr):
         """Add nodes to graph from a container ebunch.
@@ -492,10 +491,9 @@ def add_edges_from(self, ebunch, **attr):
         >>> elist=[('a','b'),('b','c')]
         >>> G.add_edges_from(elist)
 
-        Attributes can be added when edges are created
+        Attributes can be added when edges are created or updated after creation
 
         >>> G.add_edges_from(elist, color='green')
-
         """
         for e in ebunch:
             self.add_edge(e, **attr)

pygraphviz/tests/test_edge_attributes.py

@@ -64,3 +64,26 @@ def _test_anonymous_edges():
 }
 """.replace('\n', linesep))
     os.unlink(fname)
+
+def test_edge_attribute_update():
+    A = pgv.AGraph(strict=True)
+    A.add_edge(1,2,label='test',spam='eggs')
+    A.add_edge(1,2,label='update',spam='')
+    assert_equal(A.string().expandtabs(2),
+"""strict graph {
+  1 -- 2   [label=update];
+}
+""".replace('\n', linesep)
+)
+
+def test_edge_attribute_update_nonstrict():
+    A = pgv.AGraph(strict=False)
+    A.add_edge(1,2,label='test',spam='eggs',key='one')
+    A.add_edge(1,2,label='update',spam='',key='one')
+    assert_equal(A.string().expandtabs(2),
+"""graph {
+  1 -- 2 [key=one,
+  label=update];
+}
+""".replace('\n', linesep)
+)

pygraphviz/tests/test_graph.py

@@ -248,7 +248,7 @@ def test_graph_not_strict(self):
         A = pgv.AGraph(strict=False)
         assert_false(A.is_strict())
         A.add_node(1)
-        A.add_node(1)  # silent failure
+        A.add_node(1)  # nop
         assert_equal(A.nodes(), ['1'])
         A.add_edge(1,2)
         A.add_edge(1,2)

pygraphviz/tests/test_node_attributes.py

@@ -48,3 +48,26 @@ def test_node_attributes2():
   1  [label=""];
 }
 """.replace('\n', linesep))
+
+def test_node_attribute_update():
+    A = pgv.AGraph()
+    A.add_node(1,label='test',spam='eggs')
+    A.add_node(1,label='updated')
+    assert_equal(A.string().expandtabs(2),
+"""strict graph {
+  node [label="\\N"];
+  1  [label=updated,
+    spam=eggs];
+}
+""".replace('\n', linesep))
+
+def test_node_attribute_remove():
+    A = pgv.AGraph()
+    A.add_node(1,label='test',spam='eggs')
+    A.add_node(1,label=r'\N',spam='') # use \N to signify null label, otherwise ''
+    assert_equal(A.string().expandtabs(2),
+"""strict graph {
+  node [label="\\N"];
+  1;
+}
+""".replace('\n', linesep))