From 904727f40d29a78f6231e046dc81504bf3c7b02e Mon Sep 17 00:00:00 2001
From: braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Wed, 10 May 2023 13:44:52 +0200
Subject: [PATCH] refactor(core): allow passing single item to remove actions

Signed-off-by: braks <78412429+bcakmakoglu@users.noreply.github.com>
---
 packages/core/src/store/actions.ts | 24 ++++++++++++++++++------
 packages/core/src/types/store.ts   |  6 ++++--
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/packages/core/src/store/actions.ts b/packages/core/src/store/actions.ts
index f4a88a035..acf35aa81 100644
--- 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
index e0fac0592..4b296b959 100644
--- 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: