Merge pull request #9 from n-peugnet/nodes-map

Add a map of the graph's nodes to make accessing them easier

graphml/graphml.go

@@ -144,6 +144,8 @@ type Graph struct {
 
 	// The parent GraphML
 	parent *GraphML
+	// The map of nodes, indexed by their ID
+	nodesMap map[string]*Node
 	// The map of edges by connected nodes
 	edgesMap map[string]*Edge
 	// The default edge direction flag
@@ -261,8 +263,10 @@ func (gml *GraphML) Decode(r io.Reader) error {
 			gr.edgesMap[edgeIdentifier(e.Source, e.Target)] = e
 			e.graph = gr
 		}
-		// properly link nodes to their graph
+		// populate nodes map and link them to their graph
+		gr.nodesMap = make(map[string]*Node)
 		for _, n := range gr.Nodes {
+			gr.nodesMap[n.ID] = n
 			n.graph = gr
 		}
 	}
@@ -333,6 +337,7 @@ func (gml *GraphML) AddGraph(description string, edgeDefault EdgeDirection, attr
 		Nodes:          make([]*Node, 0),
 		Edges:          make([]*Edge, 0),
 		parent:         gml,
+		nodesMap:       make(map[string]*Node),
 		edgesMap:       make(map[string]*Edge),
 		edgesDirection: edgeDefault,
 	}
@@ -362,9 +367,18 @@ func (gr *Graph) AddNode(attributes map[string]interface{}, description string)
 	// add node
 	node.graph = gr
 	gr.Nodes = append(gr.Nodes, node)
+	gr.nodesMap[node.ID] = node
 	return node, nil
 }
 
+// GetNode method to test if node with given id exists. If node exists it will be returned, otherwise nil returned
+func (gr *Graph) GetNode(id string) *Node {
+	if node, ok := gr.nodesMap[id]; ok {
+		return node
+	}
+	return nil
+}
+
 // AddEdge adds edge to the graph which connects two its nodes with provided additional attributes and description
 func (gr *Graph) AddEdge(source, target *Node, attributes map[string]interface{}, edgeDirection EdgeDirection, description string) (edge *Edge, err error) {
 	// test if edge already exists
@@ -415,6 +429,16 @@ func (gr *Graph) GetEdge(sourceId, targetId string) *Edge {
 	return nil
 }
 
+// SourceNode method to get the source node struct. If it exists it will be returned, otherwise nil returned
+func (e *Edge) SourceNode() *Node {
+	return e.graph.GetNode(e.Source)
+}
+
+// TargetNode method to get the target node struct. If it exists it will be returned, otherwise nil returned
+func (e *Edge) TargetNode() *Node {
+	return e.graph.GetNode(e.Target)
+}
+
 // 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

@@ -355,6 +355,47 @@ func TestGraph_GetEdge(t *testing.T) {
 	assert.Nil(t, edge, "edge is not expected")
 }
 
+func TestGraph_GetNode(t *testing.T) {
+	description := "test graph"
+	gml := NewGraphML("")
+
+	graph, err := gml.AddGraph(description, EdgeDirectionDirected, nil)
+	require.NoError(t, err, "failed to add graph")
+	require.NotNil(t, graph)
+
+	// add nodes
+	description = "test node #1"
+	n1, err := graph.AddNode(nil, description)
+	require.NoError(t, err, "failed to add node: %s", description)
+
+	description = "test node #2"
+	n2, err := graph.AddNode(nil, description)
+	require.NoError(t, err, "failed to add node: %s", description)
+
+	// add edge
+	description = "test edge"
+	e1, err := graph.AddEdge(n1, n2, nil, EdgeDirectionDefault, description)
+	require.NoError(t, err, "failed to add edge: %s", description)
+
+	// check existing node
+	node := graph.GetNode(n1.ID)
+	assert.NotNil(t, node, "node is expected")
+	assert.Same(t, n1, node, "node should be the same as n1")
+
+	// check non-existing node
+	node = graph.GetNode("n42")
+	assert.Nil(t, node, "node is not expected")
+
+	// check edge nodes
+	node = e1.SourceNode()
+	assert.NotNil(t, node, "node is expected")
+	assert.Same(t, n1, node, "node should be the same as n1")
+
+	node = e1.TargetNode()
+	assert.NotNil(t, node, "node is expected")
+	assert.Same(t, n2, node, "node should be the same as n2")
+}
+
 // tests if default value of common key will be used for empty attributes
 func TestGraphML_RegisterKeyDefaultValue(t *testing.T) {
 	description := "graphml"