Wrap pragmatic dnd API with hooks

package.json

@@ -64,7 +64,7 @@
     "zip-dir": "^2.0.0"
   },
   "dependencies": {
-    "@atlaskit/pragmatic-drag-and-drop": "^1.2.1",
+    "@atlaskit/pragmatic-drag-and-drop": "^1.3.1",
     "@comfyorg/litegraph": "^0.8.1",
     "@primevue/themes": "^4.0.5",
     "@vueuse/core": "^11.0.0",

src/components/common/TreeExplorerTreeNode.vue

@@ -34,19 +34,16 @@
 </template>
 
 <script setup lang="ts">
-import { ref, onMounted, onUnmounted, inject, Ref, computed } from 'vue'
+import { ref, inject, Ref, computed } from 'vue'
 import Badge from 'primevue/badge'
-import {
-  dropTargetForElements,
-  draggable
-} from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
 import type {
   TreeExplorerDragAndDropData,
   RenderedTreeExplorerNode,
   TreeExplorerNode
 } from '@/types/treeExplorerTypes'
 import EditableText from '@/components/common/EditableText.vue'
 import { useErrorHandling } from '@/hooks/errorHooks'
+import { usePragmaticDraggable, usePragmaticDroppable } from '@/hooks/dndHooks'
 
 const props = defineProps<{
   node: RenderedTreeExplorerNode
@@ -89,54 +86,44 @@ const handleRename = errorHandling.wrapWithErrorHandlingAsync(
 )
 const container = ref<HTMLElement | null>(null)
 const canDrop = ref(false)
-const treeNodeElement = ref<HTMLElement | null>(null)
-let dropTargetCleanup = () => {}
-let draggableCleanup = () => {}
-onMounted(() => {
-  treeNodeElement.value = container.value?.closest(
-    '.p-tree-node-content'
-  ) as HTMLElement
-  if (props.node.droppable) {
-    dropTargetCleanup = dropTargetForElements({
-      element: treeNodeElement.value,
-      onDrop: async (event) => {
-        const dndData = event.source.data as TreeExplorerDragAndDropData
-        if (dndData.type === 'tree-explorer-node') {
-          await props.node.handleDrop?.(props.node, dndData)
-          canDrop.value = false
-          emit('itemDropped', props.node, dndData.data)
-        }
-      },
-      onDragEnter: (event) => {
-        const dndData = event.source.data as TreeExplorerDragAndDropData
-        if (dndData.type === 'tree-explorer-node') {
-          canDrop.value = true
-        }
-      },
-      onDragLeave: () => {
-        canDrop.value = false
+
+const treeNodeElementGetter = () =>
+  container.value?.closest('.p-tree-node-content') as HTMLElement
+
+if (props.node.draggable) {
+  usePragmaticDraggable(treeNodeElementGetter, {
+    getInitialData: () => {
+      return {
+        type: 'tree-explorer-node',
+        data: props.node
       }
-    })
-  }
+    },
+    onDragStart: () => emit('dragStart', props.node),
+    onDrop: () => emit('dragEnd', props.node)
+  })
+}
 
-  if (props.node.draggable) {
-    draggableCleanup = draggable({
-      element: treeNodeElement.value,
-      getInitialData() {
-        return {
-          type: 'tree-explorer-node',
-          data: props.node
-        }
-      },
-      onDragStart: () => emit('dragStart', props.node),
-      onDrop: () => emit('dragEnd', props.node)
-    })
-  }
-})
-onUnmounted(() => {
-  dropTargetCleanup()
-  draggableCleanup()
-})
+if (props.node.droppable) {
+  usePragmaticDroppable(treeNodeElementGetter, {
+    onDrop: async (event) => {
+      const dndData = event.source.data as TreeExplorerDragAndDropData
+      if (dndData.type === 'tree-explorer-node') {
+        await props.node.handleDrop?.(props.node, dndData)
+        canDrop.value = false
+        emit('itemDropped', props.node, dndData.data)
+      }
+    },
+    onDragEnter: (event) => {
+      const dndData = event.source.data as TreeExplorerDragAndDropData
+      if (dndData.type === 'tree-explorer-node') {
+        canDrop.value = true
+      }
+    },
+    onDragLeave: () => {
+      canDrop.value = false
+    }
+  })
+}
 </script>
 
 <style scoped>

src/components/graph/GraphCanvas.vue

@@ -22,10 +22,9 @@ import SideToolbar from '@/components/sidebar/SideToolbar.vue'
 import LiteGraphCanvasSplitterOverlay from '@/components/LiteGraphCanvasSplitterOverlay.vue'
 import NodeSearchboxPopover from '@/components/searchbox/NodeSearchBoxPopover.vue'
 import NodeTooltip from '@/components/graph/NodeTooltip.vue'
-import { ref, computed, onUnmounted, onMounted, watchEffect } from 'vue'
+import { ref, computed, onMounted, watchEffect } from 'vue'
 import { app as comfyApp } from '@/scripts/app'
 import { useSettingStore } from '@/stores/settingStore'
-import { dropTargetForElements } from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
 import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
 import { useWorkspaceStore } from '@/stores/workspaceStateStore'
 import {
@@ -47,6 +46,7 @@ import {
   useModelToNodeStore
 } from '@/stores/modelToNodeStore'
 import GraphCanvasMenu from '@/components/graph/GraphCanvasMenu.vue'
+import { usePragmaticDroppable } from '@/hooks/dndHooks'
 
 const emit = defineEmits(['ready'])
 const canvasRef = ref<HTMLCanvasElement | null>(null)
@@ -115,7 +115,60 @@ watchEffect(() => {
   canvasStore.canvas.canvas.style.cursor = 'default'
 })
 
-let dropTargetCleanup = () => {}
+usePragmaticDroppable(() => canvasRef.value, {
+  onDrop: (event) => {
+    const loc = event.location.current.input
+    const dndData = event.source.data
+
+    if (dndData.type === 'tree-explorer-node') {
+      const node = dndData.data as RenderedTreeExplorerNode
+      if (node.data instanceof ComfyNodeDefImpl) {
+        const nodeDef = node.data
+        // Add an offset on x to make sure after adding the node, the cursor
+        // is on the node (top left corner)
+        const pos = comfyApp.clientPosToCanvasPos([
+          loc.clientX - 20,
+          loc.clientY
+        ])
+        comfyApp.addNodeOnGraph(nodeDef, { pos })
+      } else if (node.data instanceof ComfyModelDef) {
+        const model = node.data
+        const pos = comfyApp.clientPosToCanvasPos([loc.clientX, loc.clientY])
+        const nodeAtPos = comfyApp.graph.getNodeOnPos(pos[0], pos[1])
+        let targetProvider: ModelNodeProvider | null = null
+        let targetGraphNode: LGraphNode | null = null
+        if (nodeAtPos) {
+          const providers = modelToNodeStore.getAllNodeProviders(
+            model.directory
+          )
+          for (const provider of providers) {
+            if (provider.nodeDef.name === nodeAtPos.comfyClass) {
+              targetGraphNode = nodeAtPos
+              targetProvider = provider
+            }
+          }
+        }
+        if (!targetGraphNode) {
+          const provider = modelToNodeStore.getNodeProvider(model.directory)
+          if (provider) {
+            targetGraphNode = comfyApp.addNodeOnGraph(provider.nodeDef, {
+              pos
+            })
+            targetProvider = provider
+          }
+        }
+        if (targetGraphNode) {
+          const widget = targetGraphNode.widgets.find(
+            (widget) => widget.name === targetProvider.key
+          )
+          if (widget) {
+            widget.value = model.file_name
+          }
+        }
+      }
+    }
+  }
+})
 
 onMounted(async () => {
   // Backward compatible
@@ -140,66 +193,6 @@ onMounted(async () => {
   window['app'] = comfyApp
   window['graph'] = comfyApp.graph
 
-  dropTargetCleanup = dropTargetForElements({
-    element: canvasRef.value,
-    onDrop: (event) => {
-      const loc = event.location.current.input
-      const dndData = event.source.data
-
-      if (dndData.type === 'tree-explorer-node') {
-        const node = dndData.data as RenderedTreeExplorerNode
-        if (node.data instanceof ComfyNodeDefImpl) {
-          const nodeDef = node.data
-          // Add an offset on x to make sure after adding the node, the cursor
-          // is on the node (top left corner)
-          const pos = comfyApp.clientPosToCanvasPos([
-            loc.clientX - 20,
-            loc.clientY
-          ])
-          comfyApp.addNodeOnGraph(nodeDef, { pos })
-        } else if (node.data instanceof ComfyModelDef) {
-          const model = node.data
-          const pos = comfyApp.clientPosToCanvasPos([loc.clientX, loc.clientY])
-          const nodeAtPos = comfyApp.graph.getNodeOnPos(pos[0], pos[1])
-          let targetProvider: ModelNodeProvider | null = null
-          let targetGraphNode: LGraphNode | null = null
-          if (nodeAtPos) {
-            const providers = modelToNodeStore.getAllNodeProviders(
-              model.directory
-            )
-            for (const provider of providers) {
-              if (provider.nodeDef.name === nodeAtPos.comfyClass) {
-                targetGraphNode = nodeAtPos
-                targetProvider = provider
-              }
-            }
-          }
-          if (!targetGraphNode) {
-            const provider = modelToNodeStore.getNodeProvider(model.directory)
-            if (provider) {
-              targetGraphNode = comfyApp.addNodeOnGraph(provider.nodeDef, {
-                pos
-              })
-              targetProvider = provider
-            }
-          }
-          if (targetGraphNode) {
-            const widget = targetGraphNode.widgets.find(
-              (widget) => widget.name === targetProvider.key
-            )
-            if (widget) {
-              widget.value = model.file_name
-            }
-          }
-        }
-      }
-    }
-  })
-
   emit('ready')
 })
-
-onUnmounted(() => {
-  dropTargetCleanup()
-})
 </script>

src/hooks/dndHooks.ts

@@ -0,0 +1,59 @@
+import { onBeforeUnmount, onMounted } from 'vue'
+import {
+  dropTargetForElements,
+  draggable
+} from '@atlaskit/pragmatic-drag-and-drop/element/adapter'
+
+export function usePragmaticDroppable(
+  dropTargetElement: HTMLElement | (() => HTMLElement),
+  options: Omit<Parameters<typeof dropTargetForElements>[0], 'element'>
+) {
+  let cleanup = () => {}
+
+  onMounted(() => {
+    const element =
+      typeof dropTargetElement === 'function'
+        ? dropTargetElement()
+        : dropTargetElement
+
+    if (!element) {
+      return
+    }
+
+    cleanup = dropTargetForElements({
+      element,
+      ...options
+    })
+  })
+
+  onBeforeUnmount(() => {
+    cleanup()
+  })
+}
+
+export function usePragmaticDraggable(
+  draggableElement: HTMLElement | (() => HTMLElement),
+  options: Omit<Parameters<typeof draggable>[0], 'element'>
+) {
+  let cleanup = () => {}
+
+  onMounted(() => {
+    const element =
+      typeof draggableElement === 'function'
+        ? draggableElement()
+        : draggableElement
+
+    if (!element) {
+      return
+    }
+
+    cleanup = draggable({
+      element,
+      ...options
+    })
+  })
+
+  onBeforeUnmount(() => {
+    cleanup()
+  })
+}