From 6a2c1b00d738a2fde0bf11d09cb3a2c3c2d4210d Mon Sep 17 00:00:00 2001
From: Braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Wed, 23 Mar 2022 15:49:00 +0100
Subject: [PATCH] feat: Enable passing components as node/edge-types

* Instead of using slots (re)-enable passing components for node/edge-types
---
 examples/CustomNode/CustomNode.vue  | 32 ++++++++++++++---------------
 src/container/VueFlow/watch.ts      | 12 ++++++++++-
 src/container/ZoomPane/ZoomPane.vue |  4 ++--
 src/store/actions.ts                |  4 ++--
 src/store/getters.ts                |  2 ++
 src/store/state.ts                  |  6 ++++--
 src/types/flow.ts                   |  5 ++++-
 src/types/store.ts                  |  8 ++++----
 8 files changed, 44 insertions(+), 29 deletions(-)

diff --git a/examples/CustomNode/CustomNode.vue b/examples/CustomNode/CustomNode.vue
index b6d0e635..43968aad 100644
--- a/examples/CustomNode/CustomNode.vue
+++ b/examples/CustomNode/CustomNode.vue
@@ -1,6 +1,6 @@
 <script lang="ts" setup>
 import ColorSelectorNode from './ColorSelectorNode.vue'
-import { ConnectionMode, Elements, isEdge, Node, Position, SnapGrid, useVueFlow, VueFlow, Controls, MiniMap } from '~/index'
+import { ConnectionMode, Controls, Elements, isEdge, MiniMap, Node, Position, SnapGrid, useVueFlow, VueFlow } from '~/index'
 
 const elements = ref<Elements>([])
 const bgColor = ref('#1A192B')
@@ -62,25 +62,23 @@ onMounted(() => {
   ]
 })
 
-const { onPaneReady } = useVueFlow()
-onPaneReady((flowInstance) => {
-  flowInstance.fitView()
-  console.log('flow loaded:', flowInstance)
+const { onPaneReady } = useVueFlow({
+  connectionMode: ConnectionMode.Loose,
+  connectionLineStyle,
+  snapToGrid: true,
+  snapGrid,
+  defaultZoom: 1.5,
+  nodeTypes: {
+    selectorNode: markRaw(ColorSelectorNode),
+  },
+})
+onPaneReady((i) => {
+  i.fitView()
+  console.log('flow loaded:', i)
 })
 </script>
 <template>
-  <VueFlow
-    v-model="elements"
-    :style="{ backgroundColor: bgColor }"
-    :connection-mode="ConnectionMode.Loose"
-    :connection-line-style="connectionLineStyle"
-    :snap-to-grid="true"
-    :snap-grid="snapGrid"
-    :default-zoom="1.5"
-  >
-    <template #node-selectorNode="props">
-      <ColorSelectorNode v-bind="props" />
-    </template>
+  <VueFlow v-model="elements" :style="{ backgroundColor: bgColor }">
     <MiniMap :node-stroke-color="nodeStroke" :node-color="nodeColor" />
     <Controls />
   </VueFlow>
diff --git a/src/container/VueFlow/watch.ts b/src/container/VueFlow/watch.ts
index 77beefa4..02308879 100644
--- a/src/container/VueFlow/watch.ts
+++ b/src/container/VueFlow/watch.ts
@@ -98,7 +98,7 @@ export default (
     )
     watch(
       () => props.paneMovable,
-      (v) => isDef(v) && (store.paneMoveable = v),
+      (v) => isDef(v) && (store.paneMovable = v),
       { immediate: isDef(props.paneMovable) },
     )
     watch(
@@ -206,5 +206,15 @@ export default (
       (v) => isDef(v) && (store.applyDefault = v),
       { immediate: isDef(props.applyDefault) },
     )
+    watch(
+      () => props.nodeTypes,
+      (v) => isDef(v) && (store.nodeTypes = v),
+      { immediate: isDef(props.nodeTypes) },
+    )
+    watch(
+      () => props.edgeTypes,
+      (v) => isDef(v) && (store.edgeTypes = v),
+      { immediate: isDef(props.edgeTypes) },
+    )
   })
 }
diff --git a/src/container/ZoomPane/ZoomPane.vue b/src/container/ZoomPane/ZoomPane.vue
index 8d76867e..ab1e26fb 100644
--- a/src/container/ZoomPane/ZoomPane.vue
+++ b/src/container/ZoomPane/ZoomPane.vue
@@ -121,7 +121,7 @@ onMounted(async () => {
     const pinchZoom = store.zoomOnPinch && event.ctrlKey
 
     // if all interactions are disabled, we prevent all zoom events
-    if (!store.paneMoveable && !store.zoomOnScroll && !store.panOnScroll && !store.zoomOnDoubleClick && !store.zoomOnPinch)
+    if (!store.paneMovable && !store.zoomOnScroll && !store.panOnScroll && !store.zoomOnDoubleClick && !store.zoomOnPinch)
       return false
 
     // during a selection we prevent all other interactions
@@ -148,7 +148,7 @@ onMounted(async () => {
     if (!store.zoomOnScroll && !store.panOnScroll && !pinchZoom && event.type === 'wheel') return false
 
     // if the pane is not movable, we prevent dragging it with mousestart or touchstart
-    if (!store.paneMoveable && (event.type === 'mousedown' || event.type === 'touchstart')) return false
+    if (!store.paneMovable && (event.type === 'mousedown' || event.type === 'touchstart')) return false
 
     // default filter for d3-zoom
     return (!event.ctrlKey || event.type === 'wheel') && !event.button
diff --git a/src/store/actions.ts b/src/store/actions.ts
index 98bc5e4c..a465a5cc 100644
--- a/src/store/actions.ts
+++ b/src/store/actions.ts
@@ -209,7 +209,7 @@ export default (state: State, getters: ComputedGetters): Actions => {
   }
 
   const setElements: Actions['setElements'] = (elements, extent) => {
-    if (!state.initialized && !elements.length) return
+    if ((!state.initialized && !elements.length) || !elements) return
     setNodes(elements.filter(isNode), extent)
     setEdges(elements.filter(isEdge))
   }
@@ -255,7 +255,7 @@ export default (state: State, getters: ComputedGetters): Actions => {
     if (typeof opts.panOnScroll !== 'undefined') state.panOnScroll = opts.panOnScroll
     if (typeof opts.panOnScrollMode !== 'undefined') state.panOnScrollMode = opts.panOnScrollMode
     if (typeof opts.panOnScrollSpeed !== 'undefined') state.panOnScrollSpeed = opts.panOnScrollSpeed
-    if (typeof opts.paneMovable !== 'undefined') state.paneMoveable = opts.paneMovable
+    if (typeof opts.paneMovable !== 'undefined') state.paneMovable = opts.paneMovable
     if (typeof opts.zoomOnScroll !== 'undefined') state.zoomOnScroll = opts.zoomOnScroll
     if (typeof opts.preventScrolling !== 'undefined') state.preventScrolling = opts.preventScrolling
     if (typeof opts.zoomOnDoubleClick !== 'undefined') state.zoomOnDoubleClick = opts.zoomOnDoubleClick
diff --git a/src/store/getters.ts b/src/store/getters.ts
index 1f8c5eee..b9d0f16a 100644
--- a/src/store/getters.ts
+++ b/src/store/getters.ts
@@ -6,6 +6,7 @@ export default (state: State): ComputedGetters => {
   const getEdgeTypes = computed(() => {
     const edgeTypes: Record<string, any> = {
       ...defaultEdgeTypes,
+      ...state.edgeTypes,
     }
     const keys = Object.keys(edgeTypes)
     state.edges?.forEach((e) => e.type && !keys.includes(e.type) && (edgeTypes[e.type] = e.type))
@@ -15,6 +16,7 @@ export default (state: State): ComputedGetters => {
   const getNodeTypes = computed(() => {
     const nodeTypes: Record<string, any> = {
       ...defaultNodeTypes,
+      ...state.nodeTypes,
     }
     const keys = Object.keys(nodeTypes)
     state.nodes?.forEach((n) => n.type && !keys.includes(n.type) && (nodeTypes[n.type] = n.type))
diff --git a/src/store/state.ts b/src/store/state.ts
index 8e53b725..5223d3ae 100644
--- a/src/store/state.ts
+++ b/src/store/state.ts
@@ -59,7 +59,7 @@ export default (opts?: FlowOptions): State => {
     panOnScroll: false,
     panOnScrollSpeed: 0.5,
     panOnScrollMode: PanOnScrollMode.Free,
-    paneMoveable: true,
+    paneMovable: true,
     edgeUpdaterRadius: 10,
     onlyRenderVisibleElements: false,
     defaultZoom: 1,
@@ -100,10 +100,12 @@ export default (opts?: FlowOptions): State => {
   } as State
 
   if (opts) {
+    if (typeof opts.nodeTypes !== 'undefined') state.nodeTypes = opts.nodeTypes
+    if (typeof opts.edgeTypes !== 'undefined') state.edgeTypes = opts.edgeTypes
     if (typeof opts.panOnScroll !== 'undefined') state.panOnScroll = opts.panOnScroll
     if (typeof opts.panOnScrollMode !== 'undefined') state.panOnScrollMode = opts.panOnScrollMode
     if (typeof opts.panOnScrollSpeed !== 'undefined') state.panOnScrollSpeed = opts.panOnScrollSpeed
-    if (typeof opts.paneMovable !== 'undefined') state.paneMoveable = opts.paneMovable
+    if (typeof opts.paneMovable !== 'undefined') state.paneMovable = opts.paneMovable
     if (typeof opts.zoomOnScroll !== 'undefined') state.zoomOnScroll = opts.zoomOnScroll
     if (typeof opts.preventScrolling !== 'undefined') state.preventScrolling = opts.preventScrolling
     if (typeof opts.zoomOnDoubleClick !== 'undefined') state.zoomOnDoubleClick = opts.zoomOnDoubleClick
diff --git a/src/types/flow.ts b/src/types/flow.ts
index 13090fed..d012eefb 100644
--- a/src/types/flow.ts
+++ b/src/types/flow.ts
@@ -5,6 +5,7 @@ import { ConnectionLineType, ConnectionMode } from './connection'
 import { KeyCode, PanOnScrollMode, UseZoomPanHelper } from './zoom'
 import { Actions, ComputedGetters, State, FlowStore } from './store'
 import { FlowHooksOn } from './hooks'
+import { DefaultEdgeTypes, DefaultNodeTypes, EdgeComponent, NodeComponent } from "~/types/components";
 
 export type FlowElement<N = any, E = any> = GraphNode<N> | GraphEdge<E>
 export type FlowElements<N = any, E = any> = FlowElement<N, E>[]
@@ -83,11 +84,13 @@ export type FlowInstance<N = any, E = N> = {
 } & UseZoomPanHelper
 
 export interface FlowProps<N = any, E = N> {
+  id?: string
   /** @deprecated use nodes / edges instead */
   modelValue?: Elements<N, E>
   nodes?: Node<N>[]
   edges?: Edge<E>[]
-  id?: string
+  edgeTypes?: { [key in keyof DefaultEdgeTypes]?: EdgeComponent } & { [key: string]: EdgeComponent }
+  nodeTypes?: { [key in keyof DefaultNodeTypes]?: NodeComponent } & { [key: string]: NodeComponent }
   connectionMode?: ConnectionMode
   connectionLineType?: ConnectionLineType
   connectionLineStyle?: CSSProperties
diff --git a/src/types/store.ts b/src/types/store.ts
index 594ea2c7..a4152c9b 100644
--- a/src/types/store.ts
+++ b/src/types/store.ts
@@ -1,6 +1,6 @@
 import { ComputedRef, CSSProperties, ToRefs } from 'vue'
 import { Dimensions, Elements, FlowElements, FlowInstance, FlowOptions, Rect, SnapGrid, Transform, XYPosition } from './flow'
-import { HandleType, EdgeComponent, NodeComponent } from './components'
+import { HandleType, EdgeComponent, NodeComponent, DefaultNodeTypes, DefaultEdgeTypes } from './components'
 import { Connection, ConnectionLineType, ConnectionMode, SetConnectionId } from './connection'
 import { Edge, GraphEdge } from './edge'
 import { GraphNode, CoordinateExtent, Node } from './node'
@@ -54,7 +54,7 @@ export interface State<N = any, E = N> extends Omit<FlowOptions<N, E>, 'id' | 'm
   nodesConnectable: boolean
   elementsSelectable: boolean
   selectNodesOnDrag: boolean
-  paneMoveable: boolean
+  paneMovable: boolean
   zoomOnScroll: boolean
   zoomOnPinch: boolean
   panOnScroll: boolean
@@ -95,8 +95,8 @@ export interface Actions<N = any, E = N> {
 }
 
 export interface Getters<N = any, E = N> {
-  getEdgeTypes: Record<string, EdgeComponent>
-  getNodeTypes: Record<string, NodeComponent>
+  getEdgeTypes: Record<keyof DefaultEdgeTypes | string, EdgeComponent>
+  getNodeTypes: Record<keyof DefaultNodeTypes | string, NodeComponent>
   getNodes: GraphNode<N>[]
   getEdges: GraphEdge<E>[]
   getNode: (id: string) => GraphNode<N> | undefined