Merge pull request #10 from n-peugnet/remove-key-attrs

Add methods to remove attributes from elements & keys from GraphML

graphml/graphml.go

@@ -304,6 +304,48 @@ func (gml *GraphML) RegisterKey(target KeyForElement, name, description string,
 	return key, nil
 }
 
+// RemoveKey removes a key from the GraphML and all the associated attributes
+// in all the target elements.
+func (gml *GraphML) RemoveKey(key *Key) error {
+	var found bool
+	var i int
+	var k *Key
+	for i, k = range gml.Keys {
+		if key == k {
+			found = true
+			break
+		}
+	}
+	if !found {
+		return errors.New("key not found")
+	}
+	gml.Keys = append(gml.Keys[:i], gml.Keys[i+1:]...)
+	delete(gml.keysById, key.ID)
+	delete(gml.keysByIdentifier, keyIdentifier(key.Name, key.Target))
+	if key.Target == KeyForAll || key.Target == KeyForGraphML {
+		gml.RemoveAttribute(key.ID)
+	}
+	if key.Target == KeyForGraphML {
+		return nil
+	}
+	for _, graph := range gml.Graphs {
+		if key.Target == KeyForAll || key.Target == KeyForGraph {
+			graph.RemoveAttribute(key.ID)
+		}
+		if key.Target == KeyForAll || key.Target == KeyForNode {
+			for _, node := range graph.Nodes {
+				node.RemoveAttribute(key.ID)
+			}
+		}
+		if key.Target == KeyForAll || key.Target == KeyForEdge {
+			for _, edge := range graph.Edges {
+				edge.RemoveAttribute(key.ID)
+			}
+		}
+	}
+	return nil
+}
+
 // GetKey looks for registered keys with specified name for a given target element. If specific target has no
 // registered key then common target (KeyForAll) will be checked next. Returns Key (either specific or common) or nil.
 func (gml *GraphML) GetKey(name string, target KeyForElement) *Key {
@@ -439,6 +481,43 @@ func (e *Edge) TargetNode() *Node {
 	return e.graph.GetNode(e.Target)
 }
 
+// RemoveAttribute removes the attribute associated with the given key ID from
+// the data of this GraphML.
+func (gml *GraphML) RemoveAttribute(key string) {
+	gml.Data = removeAttributeFromData(gml.Data, key)
+}
+
+// RemoveAttribute removes the attribute associated with the given key ID from
+// the data of this graph.
+func (gr *Graph) RemoveAttribute(key string) {
+	gr.Data = removeAttributeFromData(gr.Data, key)
+}
+
+// RemoveAttribute removes the attribute associated with the given key ID from
+// the data of this node.
+func (n *Node) RemoveAttribute(key string) {
+	n.Data = removeAttributeFromData(n.Data, key)
+}
+
+// RemoveAttribute removes the attribute associated with the given key ID from
+// the data of this edge.
+func (e *Edge) RemoveAttribute(key string) {
+	e.Data = removeAttributeFromData(e.Data, key)
+}
+
+// removeAttributeFromData removes the attribute associated with the given key ID from
+// the given data.
+func removeAttributeFromData(data []*Data, key string) []*Data {
+	var i int
+	var d *Data
+	for i, d = range data {
+		if d.Key == key {
+			return append(data[:i], data[i+1:]...)
+		}
+	}
+	return data
+}
+
 // GetAttributes return data attributes map associated with GraphML
 func (gml *GraphML) GetAttributes() (map[string]interface{}, error) {
 	return attributesForData(gml.Data, KeyForGraphML, gml)

graphml/graphml_test.go

@@ -548,6 +548,101 @@ func TestGraph_AddNode(t *testing.T) {
 	assert.Equal(t, attributes, attrs)
 }
 
+func TestGraph_RemoveKey(t *testing.T) {
+	gmlattrs := map[string]interface{}{
+		"k4": 8000,
+	}
+	gml, err := NewGraphMLWithAttributes("", gmlattrs)
+	require.NoError(t, err, "failed to create GraphML")
+	k1name := "k1"
+	k1target := KeyForAll
+	k1, err := gml.RegisterKey(k1target, k1name, "", reflect.Int, 0)
+	require.NoError(t, err, "failed to register key: %s", k1name)
+	k2name := "k2"
+	k2target := KeyForNode
+	k2, err := gml.RegisterKey(k2target, k2name, "", reflect.Int, 0)
+	require.NoError(t, err, "failed to register key: %s", k2name)
+	k3name := "k3"
+	k3target := KeyForEdge
+	k3, err := gml.RegisterKey(k3target, k3name, "", reflect.Int, 0)
+	require.NoError(t, err, "failed to register key: %s", k3name)
+	require.Len(t, gml.Keys, 4)
+
+	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")
+	n2attrs := map[string]interface{}{
+		"k1": 200,
+		"k2": 20,
+	}
+	n2, err := gr.AddNode(n2attrs, "test node 2")
+	require.NoError(t, err, "failed to add node 2")
+	e1attrs := map[string]interface{}{
+		"k1": 300,
+		"k3": 3,
+	}
+	e1, err := gr.AddEdge(n1, n2, e1attrs, EdgeDirectionDefault, "test edge")
+	require.NoError(t, err, "failed to add edge")
+
+	// try removing k1
+	err = gml.RemoveKey(k1)
+	require.NoError(t, err, "failed to remove key 1")
+	key := gml.GetKey(k1name, k1target)
+	assert.Nil(t, key)
+	attrs, _ := gr.GetAttributes()
+	assert.NotContains(t, attrs, "k1")
+	attrs, _ = n1.GetAttributes()
+	assert.NotContains(t, attrs, "k1")
+	attrs, _ = n2.GetAttributes()
+	assert.NotContains(t, attrs, "k1")
+	attrs, _ = e1.GetAttributes()
+	assert.NotContains(t, attrs, "k1")
+
+	// try removing k2
+	err = gml.RemoveKey(k2)
+	require.NoError(t, err, "failed to remove key 2")
+	key = gml.GetKey(k2name, k2target)
+	assert.Nil(t, key)
+	attrs, _ = n1.GetAttributes()
+	assert.NotContains(t, attrs, "k2")
+	attrs, _ = n2.GetAttributes()
+	assert.NotContains(t, attrs, "k2")
+
+	// try removing k3
+	err = gml.RemoveKey(k3)
+	require.NoError(t, err, "failed to remove key 3")
+	key = gml.GetKey(k3name, k3target)
+	assert.Nil(t, key)
+	attrs, _ = e1.GetAttributes()
+	assert.NotContains(t, attrs, "k3")
+
+	// try removing k4
+	k4name := "k4"
+	k4target := KeyForGraphML
+	k4 := gml.GetKey(k4name, k4target)
+	assert.NotNil(t, k4)
+	err = gml.RemoveKey(k4)
+	require.NoError(t, err, "failed to remove key 4")
+	key = gml.GetKey(k4name, k4target)
+	assert.Nil(t, key)
+	attrs, _ = gml.GetAttributes()
+	assert.NotContains(t, attrs, "k4")
+
+	// try removing k1 once again
+	err = gml.RemoveKey(k1)
+	assert.Error(t, err)
+}
+
 func TestNode_GetAttributes(t *testing.T) {
 	description := "test graph"
 	gml := NewGraphML("")