Implement NodeSlot (#476)

* Implement NodeSlot

* nit

src/LGraphCanvas.ts

@@ -67,6 +67,7 @@ import { getAllNestedItems, findFirstNode } from "./utils/collections"
 import { CanvasPointer } from "./CanvasPointer"
 import { BooleanWidget } from "./widgets/BooleanWidget"
 import { toClass } from "./utils/type"
+import { NodeInputSlot, NodeOutputSlot, type ConnectionColorContext } from "./NodeSlot"
 
 interface IShowSearchOptions {
   node_to?: LGraphNode
@@ -177,7 +178,7 @@ interface IPasteFromClipboardOptions {
  * This class is in charge of rendering one graph inside a canvas. And provides all the interaction required.
  * Valid callbacks are: onNodeSelected, onNodeDeselected, onShowNodePanel, onNodeDblClicked
  */
-export class LGraphCanvas {
+export class LGraphCanvas implements ConnectionColorContext {
   // Optimised buffers used during rendering
   static #temp = new Float32Array(4)
   static #temp_vec2 = new Float32Array(2)
@@ -4877,9 +4878,7 @@ export class LGraphCanvas {
       // input connection slots
       if (node.inputs) {
         for (let i = 0; i < node.inputs.length; i++) {
-          const slot = node.inputs[i]
-
-          const slot_type = slot.type
+          const slot = toClass(NodeInputSlot, node.inputs[i])
 
           // change opacity of incompatible slots when dragging a connection
           const isValid =
@@ -4891,15 +4890,7 @@ export class LGraphCanvas {
             : LiteGraph.NODE_TEXT_COLOR
           ctx.globalAlpha = isValid ? editor_alpha : 0.4 * editor_alpha
 
-          ctx.fillStyle =
-            slot.link != null
-              ? slot.color_on ||
-                this.default_connection_color_byType[slot_type] ||
-                this.default_connection_color.input_on
-              : slot.color_off ||
-                this.default_connection_color_byTypeOff[slot_type] ||
-                this.default_connection_color_byType[slot_type] ||
-                this.default_connection_color.input_off
+          ctx.fillStyle = slot.renderingColor(this)
 
           const pos = node.getConnectionPos(true, i, slot_pos)
           pos[0] -= node.pos[0]
@@ -4926,7 +4917,7 @@ export class LGraphCanvas {
       ctx.strokeStyle = "black"
       if (node.outputs) {
         for (let i = 0; i < node.outputs.length; i++) {
-          const slot = node.outputs[i]
+          const slot = toClass(NodeOutputSlot, node.outputs[i])
 
           const slot_type = slot.type
 
@@ -4947,16 +4938,7 @@ export class LGraphCanvas {
             max_y = pos[1] + LiteGraph.NODE_SLOT_HEIGHT * 0.5
           }
 
-          ctx.fillStyle =
-            slot.links && slot.links.length
-              ? slot.color_on ||
-                this.default_connection_color_byType[slot_type] ||
-                this.default_connection_color.output_on
-              : slot.color_off ||
-                this.default_connection_color_byTypeOff[slot_type] ||
-                this.default_connection_color_byType[slot_type] ||
-                this.default_connection_color.output_off
-
+          ctx.fillStyle = slot.renderingColor(this)
           drawSlot(ctx, slot, pos, {
             horizontal,
             low_quality,

src/LGraphNode.ts

@@ -33,6 +33,7 @@ import { type LGraphNodeConstructor, LiteGraph } from "./litegraph"
 import { isInRectangle, isInRect, snapPoint } from "./measure"
 import { LLink } from "./LLink"
 import { BooleanWidget } from "./widgets/BooleanWidget"
+import { NodeInputSlot, NodeOutputSlot } from "./NodeSlot"
 
 export type NodeId = number | string
 
@@ -1308,7 +1309,7 @@ export class LGraphNode implements Positionable, IPinnable {
     type?: ISlotType,
     extra_info?: object,
   ): INodeOutputSlot {
-    const output = { name: name, type: type, links: null }
+    const output = new NodeOutputSlot({ name: name, type: type, links: null })
     if (extra_info) {
       for (const i in extra_info) {
         output[i] = extra_info[i]
@@ -1334,7 +1335,7 @@ export class LGraphNode implements Positionable, IPinnable {
   addOutputs(array: [string, ISlotType, Record<string, unknown>][]): void {
     for (let i = 0; i < array.length; ++i) {
       const info = array[i]
-      const o = { name: info[0], type: info[1], links: null }
+      const o = new NodeOutputSlot({ name: info[0], type: info[1], links: null })
       if (array[2]) {
         for (const j in info[2]) {
           o[j] = info[2][j]
@@ -1383,7 +1384,7 @@ export class LGraphNode implements Positionable, IPinnable {
    */
   addInput(name: string, type: ISlotType, extra_info?: object): INodeInputSlot {
     type = type || 0
-    const input: INodeInputSlot = { name: name, type: type, link: null }
+    const input: INodeInputSlot = new NodeInputSlot({ name: name, type: type, link: null })
     if (extra_info) {
       for (const i in extra_info) {
         input[i] = extra_info[i]
@@ -1408,7 +1409,7 @@ export class LGraphNode implements Positionable, IPinnable {
   addInputs(array: [string, ISlotType, Record<string, unknown>][]): void {
     for (let i = 0; i < array.length; ++i) {
       const info = array[i]
-      const o: INodeInputSlot = { name: info[0], type: info[1], link: null }
+      const o: INodeInputSlot = new NodeInputSlot({ name: info[0], type: info[1], link: null })
       // TODO: Checking the wrong variable here - confirm no downstream consumers, then remove.
       if (array[2]) {
         for (const j in info[2]) {

src/NodeSlot.ts

@@ -0,0 +1,95 @@
+import type { CanvasColour, Dictionary, INodeInputSlot, INodeOutputSlot, INodeSlot, ISlotType, Point } from "./interfaces"
+import type { IWidget } from "./types/widgets"
+import type { LinkDirection, RenderShape } from "./types/globalEnums"
+import type { LinkId } from "./LLink"
+
+export interface ConnectionColorContext {
+  default_connection_color: {
+    input_off: string
+    input_on: string
+    output_off: string
+    output_on: string
+  }
+  default_connection_color_byType: Dictionary<CanvasColour>
+  default_connection_color_byTypeOff: Dictionary<CanvasColour>
+}
+
+export abstract class NodeSlot implements INodeSlot {
+  name: string
+  localized_name?: string
+  label?: string
+  type: ISlotType
+  dir?: LinkDirection
+  removable?: boolean
+  shape?: RenderShape
+  color_off?: CanvasColour
+  color_on?: CanvasColour
+  locked?: boolean
+  nameLocked?: boolean
+  pos?: Point
+  widget?: IWidget
+
+  constructor(slot: INodeSlot) {
+    Object.assign(this, slot)
+    this.name = slot.name
+    this.type = slot.type
+  }
+
+  /**
+   * The label to display in the UI.
+   */
+  get displayLabel(): string {
+    return this.label || this.localized_name || this.name || ""
+  }
+
+  abstract isConnected(): boolean
+
+  connectedColor(context: ConnectionColorContext): CanvasColour {
+    return this.color_on ||
+      context.default_connection_color_byType[this.type] ||
+      context.default_connection_color.output_on
+  }
+
+  disconnectedColor(context: ConnectionColorContext): CanvasColour {
+    return this.color_off ||
+      context.default_connection_color_byTypeOff[this.type] ||
+      context.default_connection_color_byType[this.type] ||
+      context.default_connection_color.output_off
+  }
+
+  renderingColor(context: ConnectionColorContext): CanvasColour {
+    return this.isConnected()
+      ? this.connectedColor(context)
+      : this.disconnectedColor(context)
+  }
+}
+
+export class NodeInputSlot extends NodeSlot implements INodeInputSlot {
+  link: LinkId | null
+
+  constructor(slot: INodeInputSlot) {
+    super(slot)
+    this.link = slot.link
+  }
+
+  override isConnected(): boolean {
+    return this.link != null
+  }
+}
+
+export class NodeOutputSlot extends NodeSlot implements INodeOutputSlot {
+  links: LinkId[] | null
+  _data?: unknown
+  slot_index?: number
+
+  constructor(slot: INodeOutputSlot) {
+    super(slot)
+    this.links = slot.links
+    this._data = slot._data
+    this.slot_index = slot.slot_index
+  }
+
+  override isConnected(): boolean {
+    return this.links != null && this.links.length > 0
+  }
+}