diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/Graph.test.ts b/invokeai/frontend/web/src/features/nodes/util/graph/Graph.test.ts index b11e16545f..412d886720 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/Graph.test.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/Graph.test.ts @@ -123,6 +123,34 @@ describe('Graph', () => { }); }); + describe('deleteNode', () => { + it('should delete the node with the provided id', () => { + const g = new Graph(); + const n1 = g.addNode({ + id: 'n1', + type: 'add', + }); + const n2 = g.addNode({ + id: 'n2', + type: 'add', + }); + const n3 = g.addNode({ + id: 'n3', + type: 'add', + }); + g.addEdge(n1, 'value', n2, 'a'); + g.addEdge(n2, 'value', n3, 'a'); + // This edge should not be deleted bc it doesn't touch n2 + g.addEdge(n1, 'value', n3, 'a'); + g.deleteNode(n2.id); + expect(g.hasNode(n1.id)).toBe(true); + expect(g.hasNode(n2.id)).toBe(false); + expect(g.hasNode(n3.id)).toBe(true); + // Should delete edges to and from the node + expect(g.getEdges().length).toBe(1); + }); + }); + describe('hasNode', () => { const g = new Graph(); g.addNode({ @@ -160,6 +188,27 @@ describe('Graph', () => { }); }); + describe('getEdges', () => { + it('should get all edges in the graph', () => { + const g = new Graph(); + const n1 = g.addNode({ + id: 'n1', + type: 'add', + }); + const n2 = g.addNode({ + id: 'n2', + type: 'add', + }); + const n3 = g.addNode({ + id: 'n3', + type: 'add', + }); + const e1 = g.addEdge(n1, 'value', n2, 'a'); + const e2 = g.addEdge(n2, 'value', n3, 'a'); + expect(g.getEdges()).toEqual([e1, e2]); + }); + }); + describe('hasEdge', () => { const g = new Graph(); const add: Invocation<'add'> = { diff --git a/invokeai/frontend/web/src/features/nodes/util/graph/Graph.ts b/invokeai/frontend/web/src/features/nodes/util/graph/Graph.ts index b578c5b40a..165e04ffe4 100644 --- a/invokeai/frontend/web/src/features/nodes/util/graph/Graph.ts +++ b/invokeai/frontend/web/src/features/nodes/util/graph/Graph.ts @@ -68,6 +68,19 @@ export class Graph { return node; } + /** + * Deletes a node from the graph. All edges to and from the node are also deleted. + * @param id The id of the node to delete. + */ + deleteNode(id: string): void { + const node = this._graph.nodes[id]; + if (node) { + this.deleteEdgesFrom(node); + this.deleteEdgesTo(node); + delete this._graph.nodes[id]; + } + } + /** * Check if a node exists in the graph. * @param id The id of the node to check. @@ -182,6 +195,14 @@ export class Graph { return edge; } + /** + * Get all edges in the graph. + * @returns The edges. + */ + getEdges(): Edge[] { + return this._graph.edges; + } + /** * Check if a graph has an edge. * Provide the from and to node types as generics to get type hints for from and to field names.