From 7fa9203ccbe290dc98fc00077fc731e251b38f7f Mon Sep 17 00:00:00 2001
From: braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Wed, 1 Feb 2023 23:31:54 +0100
Subject: [PATCH] feat(core): allow setting node/edge type with generic

Signed-off-by: braks <78412429+bcakmakoglu@users.noreply.github.com>
---
 packages/core/src/types/edge.ts | 17 ++++++++++++-----
 packages/core/src/types/node.ts | 12 ++++++++----
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/packages/core/src/types/edge.ts b/packages/core/src/types/edge.ts
index 256dd861..41645fc7 100644
--- a/packages/core/src/types/edge.ts
+++ b/packages/core/src/types/edge.ts
@@ -147,20 +147,29 @@ export interface EdgePositions {
 }
 
 /** Internal edge type */
-export type GraphEdge<Data = ElementData, CustomEvents extends Record<string, CustomEvent> = any> = Edge<Data, CustomEvents> & {
+export type GraphEdge<
+  Data = ElementData,
+  CustomEvents extends Record<string, CustomEvent> = any,
+  Type extends string = string,
+> = Edge<Data, CustomEvents> & {
   selected?: boolean
   sourceNode: GraphNode
   targetNode: GraphNode
   data: Data
   events: Partial<EdgeEventsHandler<CustomEvents>>
+  type: Type
 } & EdgePositions
 
 /** these props are passed to edge components */
-export interface EdgeProps<Data = ElementData, CustomEvents = {}> extends EdgeLabelOptions, EdgePositions {
+export interface EdgeProps<Data = ElementData, CustomEvents = {}, Type extends string = string>
+  extends EdgeLabelOptions,
+    EdgePositions {
   id: string
   sourceNode: GraphNode
   targetNode: GraphNode
-  type?: keyof DefaultEdgeTypes | string
+  source: string
+  target: string
+  type?: Type
   label?: string | VNode | Component<EdgeTextProps> | Object
   style?: CSSProperties
   selected?: boolean
@@ -168,8 +177,6 @@ export interface EdgeProps<Data = ElementData, CustomEvents = {}> extends EdgeLa
   targetPosition: Position
   sourceHandleId?: string
   targetHandleId?: string
-  source: string
-  target: string
   animated?: boolean
   updatable?: boolean
   markerStart: string
diff --git a/packages/core/src/types/node.ts b/packages/core/src/types/node.ts
index f98756a7..c9b3a434 100644
--- a/packages/core/src/types/node.ts
+++ b/packages/core/src/types/node.ts
@@ -86,8 +86,11 @@ export interface Node<Data = ElementData, CustomEvents extends Record<string, Cu
   ariaLabel?: string
 }
 
-export interface GraphNode<Data = ElementData, CustomEvents extends Record<string, CustomEvent> = any>
-  extends Node<Data, CustomEvents> {
+export interface GraphNode<
+  Data = ElementData,
+  CustomEvents extends Record<string, CustomEvent> = any,
+  Type extends string = string,
+> extends Node<Data, CustomEvents> {
   /** absolute position in relation to parent elements + z-index */
   computedPosition: XYZPosition
   handleBounds: NodeHandleBounds
@@ -100,14 +103,15 @@ export interface GraphNode<Data = ElementData, CustomEvents extends Record<strin
   initialized: boolean
   data: Data
   events: Partial<NodeEventsHandler<CustomEvents>>
+  type: Type
 }
 
 /** these props are passed to node components */
-export interface NodeProps<Data = ElementData, CustomEvents = {}> {
+export interface NodeProps<Data = ElementData, CustomEvents = {}, Type extends string = keyof DefaultNodeTypes> {
   /** unique node id */
   id: string
   /** node type */
-  type: keyof DefaultNodeTypes | string
+  type?: Type
   /** is node selected */
   selected: boolean
   /** can node handles be connected */