Add attribute setters for GraphML, Graph, Node and Edge

graphml/graphml.go

@@ -526,6 +526,50 @@ func removeAttributeFromData(data []*Data, key string) []*Data {
 	return data
 }
 
+// SetAttribute sets the value of the attribute associated with the given key ID
+// in the data of this GraphML.
+func (gml *GraphML) SetAttribute(key string, val interface{}) (err error) {
+	gml.Data, err = gml.setAttributeForData(gml.Data, KeyForGraphML, key, val)
+	return
+}
+
+// SetAttribute sets the value of the attribute associated with the given key ID
+// in the data of this graph.
+func (gr *Graph) SetAttribute(key string, val interface{}) (err error) {
+	gr.Data, err = gr.parent.setAttributeForData(gr.Data, KeyForGraph, key, val)
+	return
+}
+
+// SetAttribute sets the value of the attribute associated with the given key ID
+// in the data of this node.
+func (n *Node) SetAttribute(key string, val interface{}) (err error) {
+	n.Data, err = n.graph.parent.setAttributeForData(n.Data, KeyForNode, key, val)
+	return
+}
+
+// SetAttribute sets the value of the attribute associated with the given key ID
+// in the data of this edge.
+func (e *Edge) SetAttribute(key string, val interface{}) (err error) {
+	e.Data, err = e.graph.parent.setAttributeForData(e.Data, KeyForEdge, key, val)
+	return
+}
+
+// setAttributeForData sets the value of the attribute associated with
+// the given key ID in the given data.
+func (gml *GraphML) setAttributeForData(data []*Data, target KeyForElement, key string, val interface{}) ([]*Data, error) {
+	newData, err := gml.createDataAttribute(val, key, target)
+	if err != nil {
+		return data, err
+	}
+	for i, d := range data {
+		if d.Key == newData.Key {
+			data[i] = newData
+			return data, nil
+		}
+	}
+	return append(data, newData), nil
+}
+
 // GetAttributes return data attributes map associated with GraphML
 func (gml *GraphML) GetAttributes() (map[string]interface{}, error) {
 	return attributesForData(gml.Data, KeyForGraphML, gml)
@@ -597,7 +641,7 @@ func (gml *GraphML) addKey(key *Key) {
 // Creates data-functions from given attributes and appends definitions of created functions to the provided data list.
 func (gml *GraphML) createDataAttributes(attributes map[string]interface{}, target KeyForElement) (data []*Data, err error) {
 	// make sure that attributes are sorted in predictable order
-	names := make([]string, 0)
+	names := make([]string, 0, len(attributes))
 	for key := range attributes {
 		names = append(names, key)
 	}
@@ -606,16 +650,8 @@ func (gml *GraphML) createDataAttributes(attributes map[string]interface{}, targ
 	data = make([]*Data, len(attributes))
 	count := 0
 	for _, key := range names {
-		keyFunc := gml.GetKey(key, target)
 		val := attributes[key]
-		if keyFunc == nil {
-			// register new Key
-			if keyFunc, err = gml.RegisterKey(target, key, "", reflect.TypeOf(val).Kind(), nil); err != nil {
-				// failed
-				return nil, err
-			}
-		}
-		if d, err := createDataWithKey(val, keyFunc); err != nil {
+		if d, err := gml.createDataAttribute(val, key, target); err != nil {
 			// failed
 			return nil, err
 		} else {
@@ -626,6 +662,20 @@ func (gml *GraphML) createDataAttributes(attributes map[string]interface{}, targ
 	return data, nil
 }
 
+// createDataAttribute creates a single data object with given value, key name and target.
+// If there is no key with this name and target, a new one is registered.
+func (gml *GraphML) createDataAttribute(value interface{}, key string, target KeyForElement) (data *Data, err error) {
+	keyFunc := gml.GetKey(key, target)
+	if keyFunc == nil {
+		// register new Key
+		if keyFunc, err = gml.RegisterKey(target, key, "", reflect.TypeOf(value).Kind(), nil); err != nil {
+			// failed
+			return nil, err
+		}
+	}
+	return createDataWithKey(value, keyFunc)
+}
+
 // Creates data object with specified name, value and for provided Key
 func createDataWithKey(value interface{}, key *Key) (data *Data, err error) {
 	data = &Data{

graphml/graphml_test.go

@@ -699,6 +699,83 @@ func TestGraph_RemoveKeyByName(t *testing.T) {
 	assert.Error(t, err, "key no found")
 }
 
+func TestGraph_SetAttribute(t *testing.T) {
+	gmlattrs := map[string]interface{}{
+		"k4": 8000,
+	}
+	gml, err := NewGraphMLWithAttributes("", gmlattrs)
+	require.NoError(t, err, "failed to create GraphML")
+	k1name := "k1"
+	_, err = gml.RegisterKey(KeyForAll, "k1", "", reflect.Int, nil)
+	require.NoError(t, err, "failed to register key: %s", k1name)
+	k2name := "k2"
+	_, err = gml.RegisterKey(KeyForNode, k2name, "", reflect.Int, 0)
+	require.NoError(t, err, "failed to register key: %s", k2name)
+
+	grattrs := map[string]interface{}{
+		"k1": 999,
+	}
+	gr, err := gml.AddGraph("test graph", EdgeDirectionDirected, grattrs)
+	require.NoError(t, err, "failed to add graph")
+
+	// add elements
+	n1attrs := map[string]interface{}{
+		"k1": 100,
+		"k2": 10,
+	}
+	n1, err := gr.AddNode(n1attrs, "test node 1")
+	require.NoError(t, err, "failed to add node 1")
+	n2, err := gr.AddNode(map[string]interface{}{}, "test node 2")
+	require.NoError(t, err, "failed to add node 2")
+	e1, err := gr.AddEdge(n1, n2, map[string]interface{}{}, EdgeDirectionDefault, "test edge")
+	require.NoError(t, err, "failed to add edge")
+	require.Len(t, gml.Keys, 3)
+
+	// try setting registered attribute, non existing for GraphML
+	err = gml.SetAttribute(k1name, 42)
+	require.NoError(t, err, "failed to set key k1")
+	attrs, _ := gml.GetAttributes()
+	assert.Equal(t, map[string]interface{}{
+		"k4": 8000,
+		"k1": 42,
+	}, attrs)
+	require.Len(t, gml.Keys, 3)
+
+	// try setting invalid attribute
+	err = gml.SetAttribute("invalid", make(chan bool))
+	require.Error(t, err)
+	require.Len(t, gml.Keys, 3)
+
+	// try setting non existing attribute for Graph
+	err = gr.SetAttribute("test", 120)
+	require.NoError(t, err, "failed to set key test")
+	attrs, _ = gr.GetAttributes()
+	assert.Equal(t, map[string]interface{}{
+		"k1": 999,
+		"test": 120,
+	}, attrs)
+	require.Len(t, gml.Keys, 4)
+
+	// try setting existing attribute for Node
+	err = n1.SetAttribute(k2name, 20)
+	require.NoError(t, err, "failed to set key k2")
+	attrs, _ = n1.GetAttributes()
+	assert.Equal(t, map[string]interface{}{
+		"k1": 100,
+		"k2": 20,
+	}, attrs)
+	require.Len(t, gml.Keys, 4)
+
+	// try setting attribute non existing for Edge with no attributes
+	err = e1.SetAttribute("test", 11)
+	require.NoError(t, err, "failed to set key test")
+	attrs, _ = e1.GetAttributes()
+	assert.Equal(t, map[string]interface{}{
+		"test": 11,
+	}, attrs)
+	require.Len(t, gml.Keys, 5)
+}
+
 func TestNode_GetAttributes(t *testing.T) {
 	description := "test graph"
 	gml := NewGraphML("")