refactor(core): allow passing single item to remove actions

Signed-off-by: braks <78412429+bcakmakoglu@users.noreply.github.com>
bcakmakoglu · May 11, 2023 · 904727f · 904727f
1 parent ecb9b54
commit 904727f
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/packages/core/src/store/actions.ts b/packages/core/src/store/actions.ts
@@ -490,12 +490,18 @@ export function useActions(
   }
 
   const removeNodes: Actions['removeNodes'] = (nodes, removeConnectedEdges = true) => {
-    const curr = nodes instanceof Function ? nodes(state.nodes) : nodes
+    let nodesToRemove = nodes instanceof Function ? nodes(state.nodes) : nodes
+    nodesToRemove = Array.isArray(nodesToRemove) ? nodesToRemove : [nodesToRemove]
+
     const nodeChanges: NodeRemoveChange[] = []
     const edgeChanges: EdgeRemoveChange[] = []
 
-    curr.forEach((item) => {
-      const currNode = typeof item === 'string' ? findNode(item)! : item
+    nodesToRemove.forEach((item) => {
+      const currNode = typeof item === 'string' ? findNode(item) : item
+
+      if (!currNode) {
+        return
+      }
 
       if (isDef(currNode.deletable) && !currNode.deletable) {
         return
@@ -525,11 +531,17 @@ export function useActions(
   }
 
   const removeEdges: Actions['removeEdges'] = (edges) => {
-    const curr = edges instanceof Function ? edges(state.edges) : edges
+    let edgesToRemove = edges instanceof Function ? edges(state.edges) : edges
+    edgesToRemove = Array.isArray(edgesToRemove) ? edgesToRemove : [edgesToRemove]
+
     const changes: EdgeRemoveChange[] = []
 
-    curr.forEach((item) => {
-      const currEdge = typeof item === 'string' ? findEdge(item)! : item
+    edgesToRemove.forEach((item) => {
+      const currEdge = typeof item === 'string' ? findEdge(item) : item
+
+      if (!currEdge) {
+        return
+      }
 
       if (isDef(currEdge.deletable) && !currEdge.deletable) {
         return

diff --git a/packages/core/src/types/store.ts b/packages/core/src/types/store.ts
@@ -153,11 +153,13 @@ export type SetEdges = (edges: Edge[] | ((edges: GraphEdge[]) => Edge[])) => voi
 export type AddNodes = (nodes: Node | Node[] | ((nodes: GraphNode[]) => Node | Node[])) => void
 
 export type RemoveNodes = (
-  nodes: (Node[] | string[]) | ((nodes: GraphNode[]) => Node[] | string[]),
+  nodes: (string | Node) | (Node | string)[] | ((nodes: GraphNode[]) => (string | Node) | (Node | string)[]),
   removeConnectedEdges?: boolean,
 ) => void
 
-export type RemoveEdges = (edges: (Edge[] | string[]) | ((edges: GraphEdge[]) => Edge[] | string[])) => void
+export type RemoveEdges = (
+  edges: (string | Edge) | (Edge | string)[] | ((edges: GraphEdge[]) => (string | Edge) | (Edge | string)[]),
+) => void
 
 export type AddEdges = (
   edgesOrConnections: