[Refactor] Move drawSlot to NodeSlot (#477)

* [Refactor] Move drawSlot to NodeSlot

* nit

src/LGraphCanvas.ts

@@ -57,7 +57,7 @@ import {
   isInRect,
   snapPoint,
 } from "./measure"
-import { drawSlot, LabelPosition, strokeShape } from "./draw"
+import { LabelPosition, strokeShape } from "./draw"
 import { DragAndScale } from "./DragAndScale"
 import { LinkReleaseContextExtended, LiteGraph, clamp } from "./litegraph"
 import { stringOrEmpty, stringOrNull } from "./strings"
@@ -4857,7 +4857,6 @@ export class LGraphCanvas implements ConnectionColorContext {
     node.onDrawForeground?.(ctx, this, this.canvas)
 
     // connection slots
-    ctx.textAlign = horizontal ? "center" : "left"
     ctx.font = this.inner_text_font
 
     const render_text = !low_quality
@@ -4868,13 +4867,12 @@ export class LGraphCanvas implements ConnectionColorContext {
 
     const out_slot = this.connecting_links?.[0]?.output
     const in_slot = this.connecting_links?.[0]?.input
-    ctx.lineWidth = 1
 
     let max_y = 0
     const slot_pos = new Float32Array(2) // to reuse
 
     // render inputs and outputs
-    if (!node.flags.collapsed) {
+    if (!node.collapsed) {
       // input connection slots
       if (node.inputs) {
         for (let i = 0; i < node.inputs.length; i++) {
@@ -4890,8 +4888,6 @@ export class LGraphCanvas implements ConnectionColorContext {
             : LiteGraph.NODE_TEXT_COLOR
           ctx.globalAlpha = isValid ? editor_alpha : 0.4 * editor_alpha
 
-          ctx.fillStyle = slot.renderingColor(this)
-
           const pos = node.getConnectionPos(true, i, slot_pos)
           pos[0] -= node.pos[0]
           pos[1] -= node.pos[1]
@@ -4899,22 +4895,20 @@ export class LGraphCanvas implements ConnectionColorContext {
             max_y = pos[1] + LiteGraph.NODE_SLOT_HEIGHT * 0.5
           }
 
-          drawSlot(ctx, slot, pos, {
+          slot.draw(ctx, {
+            pos,
+            colorContext: this,
+            labelColor: label_color,
+            labelPosition: LabelPosition.Right,
             horizontal,
-            low_quality,
-            render_text,
-            label_color,
-            label_position: LabelPosition.Right,
-            // Input slot is not stroked.
-            do_stroke: false,
+            lowQuality: low_quality,
+            renderText: render_text,
             highlight,
           })
         }
       }
 
       // output connection slots
-      ctx.textAlign = horizontal ? "center" : "right"
-      ctx.strokeStyle = "black"
       if (node.outputs) {
         for (let i = 0; i < node.outputs.length; i++) {
           const slot = toClass(NodeOutputSlot, node.outputs[i])
@@ -4938,14 +4932,14 @@ export class LGraphCanvas implements ConnectionColorContext {
             max_y = pos[1] + LiteGraph.NODE_SLOT_HEIGHT * 0.5
           }
 
-          ctx.fillStyle = slot.renderingColor(this)
-          drawSlot(ctx, slot, pos, {
+          slot.draw(ctx, {
+            pos,
+            colorContext: this,
+            labelColor: label_color,
+            labelPosition: LabelPosition.Left,
             horizontal,
-            low_quality,
-            render_text,
-            label_color,
-            label_position: LabelPosition.Left,
-            do_stroke: true,
+            lowQuality: low_quality,
+            renderText: render_text,
             highlight,
           })
         }
@@ -5905,11 +5899,12 @@ export class LGraphCanvas implements ConnectionColorContext {
       const outline_color = w.advanced ? LiteGraph.WIDGET_ADVANCED_OUTLINE_COLOR : LiteGraph.WIDGET_OUTLINE_COLOR
 
       if (w === this.link_over_widget) {
-        ctx.fillStyle = this.default_connection_color_byType[this.link_over_widget_type] ||
-          this.default_connection_color.input_on
-
         // Manually draw a slot next to the widget simulating an input
-        drawSlot(ctx, {}, [10, y + 10], {})
+        new NodeInputSlot({
+          name: "",
+          type: this.link_over_widget_type,
+          link: 0,
+        }).draw(ctx, { pos: [10, y + 10], colorContext: this })
       }
 
       w.last_y = y

src/NodeSlot.ts

@@ -1,7 +1,8 @@
 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"
+import { LinkDirection, RenderShape } from "./types/globalEnums"
+import { LabelPosition, SlotShape, SlotType } from "./draw"
 
 export interface ConnectionColorContext {
   default_connection_color: {
@@ -14,6 +15,18 @@ export interface ConnectionColorContext {
   default_connection_color_byTypeOff: Dictionary<CanvasColour>
 }
 
+interface IDrawOptions {
+  pos: Point
+  colorContext: ConnectionColorContext
+  labelColor?: string
+  labelPosition?: LabelPosition
+  horizontal?: boolean
+  lowQuality?: boolean
+  renderText?: boolean
+  doStroke?: boolean
+  highlight?: boolean
+}
+
 export abstract class NodeSlot implements INodeSlot {
   name: string
   localized_name?: string
@@ -38,7 +51,7 @@ export abstract class NodeSlot implements INodeSlot {
   /**
    * The label to display in the UI.
    */
-  get displayLabel(): string {
+  get renderingLabel(): string {
     return this.label || this.localized_name || this.name || ""
   }
 
@@ -62,6 +75,117 @@ export abstract class NodeSlot implements INodeSlot {
       ? this.connectedColor(context)
       : this.disconnectedColor(context)
   }
+
+  draw(
+    ctx: CanvasRenderingContext2D,
+    options: IDrawOptions,
+  ) {
+    const {
+      pos,
+      colorContext,
+      labelColor = "#AAA",
+      labelPosition = LabelPosition.Right,
+      horizontal = false,
+      lowQuality = false,
+      renderText = true,
+      highlight = false,
+      doStroke: _doStroke = false,
+    } = options
+
+    // Save the current fillStyle and strokeStyle
+    const originalFillStyle = ctx.fillStyle
+    const originalStrokeStyle = ctx.strokeStyle
+    const originalLineWidth = ctx.lineWidth
+
+    const slot_type = this.type
+    const slot_shape = (
+      slot_type === SlotType.Array ? SlotShape.Grid : this.shape
+    ) as SlotShape
+
+    ctx.beginPath()
+    let doStroke = _doStroke
+    let doFill = true
+
+    ctx.fillStyle = this.renderingColor(colorContext)
+    ctx.lineWidth = 1
+    if (slot_type === SlotType.Event || slot_shape === SlotShape.Box) {
+      if (horizontal) {
+        ctx.rect(pos[0] - 5 + 0.5, pos[1] - 8 + 0.5, 10, 14)
+      } else {
+        ctx.rect(pos[0] - 6 + 0.5, pos[1] - 5 + 0.5, 14, 10)
+      }
+    } else if (slot_shape === SlotShape.Arrow) {
+      ctx.moveTo(pos[0] + 8, pos[1] + 0.5)
+      ctx.lineTo(pos[0] - 4, pos[1] + 6 + 0.5)
+      ctx.lineTo(pos[0] - 4, pos[1] - 6 + 0.5)
+      ctx.closePath()
+    } else if (slot_shape === SlotShape.Grid) {
+      const gridSize = 3
+      const cellSize = 2
+      const spacing = 3
+
+      for (let x = 0; x < gridSize; x++) {
+        for (let y = 0; y < gridSize; y++) {
+          ctx.rect(
+            pos[0] - 4 + x * spacing,
+            pos[1] - 4 + y * spacing,
+            cellSize,
+            cellSize,
+          )
+        }
+      }
+      doStroke = false
+    } else {
+      // Default rendering for circle, hollow circle.
+      if (lowQuality) {
+        ctx.rect(pos[0] - 4, pos[1] - 4, 8, 8)
+      } else {
+        let radius: number
+        if (slot_shape === SlotShape.HollowCircle) {
+          doFill = false
+          doStroke = true
+          ctx.lineWidth = 3
+          ctx.strokeStyle = ctx.fillStyle
+          radius = highlight ? 4 : 3
+        } else {
+          // Normal circle
+          radius = highlight ? 5 : 4
+        }
+        ctx.arc(pos[0], pos[1], radius, 0, Math.PI * 2)
+      }
+    }
+
+    if (doFill) ctx.fill()
+    if (!lowQuality && doStroke) ctx.stroke()
+
+    // render slot label
+    if (renderText) {
+      const text = this.renderingLabel
+      if (text) {
+        // TODO: Finish impl.  Highlight text on mouseover unless we're connecting links.
+        ctx.fillStyle = labelColor
+
+        if (labelPosition === LabelPosition.Right) {
+          if (horizontal || this.dir == LinkDirection.UP) {
+            ctx.fillText(text, pos[0], pos[1] - 10)
+          } else {
+            ctx.fillText(text, pos[0] + 10, pos[1] + 5)
+          }
+        } else {
+          if (horizontal || this.dir == LinkDirection.DOWN) {
+            ctx.fillText(text, pos[0], pos[1] - 8)
+          } else {
+            ctx.fillText(text, pos[0] - 10, pos[1] + 5)
+          }
+        }
+      }
+    }
+
+    // Restore the original fillStyle and strokeStyle
+    ctx.fillStyle = originalFillStyle
+    ctx.strokeStyle = originalStrokeStyle
+    ctx.lineWidth = originalLineWidth
+  }
 }
 
 export class NodeInputSlot extends NodeSlot implements INodeInputSlot {
@@ -75,6 +199,18 @@ export class NodeInputSlot extends NodeSlot implements INodeInputSlot {
   override isConnected(): boolean {
     return this.link != null
   }
+
+  override draw(ctx: CanvasRenderingContext2D, options: Omit<IDrawOptions, "doStroke">) {
+    const originalTextAlign = ctx.textAlign
+    ctx.textAlign = options.horizontal ? "center" : "left"
+
+    super.draw(ctx, {
+      ...options,
+      doStroke: false,
+    })
+
+    ctx.textAlign = originalTextAlign
+  }
 }
 
 export class NodeOutputSlot extends NodeSlot implements INodeOutputSlot {
@@ -92,4 +228,19 @@ export class NodeOutputSlot extends NodeSlot implements INodeOutputSlot {
   override isConnected(): boolean {
     return this.links != null && this.links.length > 0
   }
+
+  override draw(ctx: CanvasRenderingContext2D, options: Omit<IDrawOptions, "doStroke">) {
+    const originalTextAlign = ctx.textAlign
+    const originalStrokeStyle = ctx.strokeStyle
+    ctx.textAlign = options.horizontal ? "center" : "right"
+    ctx.strokeStyle = "black"
+
+    super.draw(ctx, {
+      ...options,
+      doStroke: true,
+    })
+
+    ctx.textAlign = originalTextAlign
+    ctx.strokeStyle = originalStrokeStyle
+  }
 }

src/draw.ts

@@ -29,121 +29,6 @@ export enum LabelPosition {
   Right = "right",
 }
 
-export function drawSlot(
-  ctx: CanvasRenderingContext2D,
-  slot: Partial<INodeSlot>,
-  pos: Vector2,
-  {
-    label_color = "#AAA",
-    label_position = LabelPosition.Right,
-    horizontal = false,
-    low_quality = false,
-    render_text = true,
-    do_stroke = false,
-    highlight = false,
-  }: {
-    label_color?: string
-    label_position?: LabelPosition
-    horizontal?: boolean
-    low_quality?: boolean
-    render_text?: boolean
-    do_stroke?: boolean
-    highlight?: boolean
-  } = {},
-) {
-  // Save the current fillStyle and strokeStyle
-  const originalFillStyle = ctx.fillStyle
-  const originalStrokeStyle = ctx.strokeStyle
-  const originalLineWidth = ctx.lineWidth
-
-  const slot_type = slot.type as SlotType
-  const slot_shape = (
-    slot_type === SlotType.Array ? SlotShape.Grid : slot.shape
-  ) as SlotShape
-
-  ctx.beginPath()
-  let doStroke = do_stroke
-  let doFill = true
-
-  if (slot_type === SlotType.Event || slot_shape === SlotShape.Box) {
-    if (horizontal) {
-      ctx.rect(pos[0] - 5 + 0.5, pos[1] - 8 + 0.5, 10, 14)
-    } else {
-      ctx.rect(pos[0] - 6 + 0.5, pos[1] - 5 + 0.5, 14, 10)
-    }
-  } else if (slot_shape === SlotShape.Arrow) {
-    ctx.moveTo(pos[0] + 8, pos[1] + 0.5)
-    ctx.lineTo(pos[0] - 4, pos[1] + 6 + 0.5)
-    ctx.lineTo(pos[0] - 4, pos[1] - 6 + 0.5)
-    ctx.closePath()
-  } else if (slot_shape === SlotShape.Grid) {
-    const gridSize = 3
-    const cellSize = 2
-    const spacing = 3
-
-    for (let x = 0; x < gridSize; x++) {
-      for (let y = 0; y < gridSize; y++) {
-        ctx.rect(
-          pos[0] - 4 + x * spacing,
-          pos[1] - 4 + y * spacing,
-          cellSize,
-          cellSize,
-        )
-      }
-    }
-    doStroke = false
-  } else {
-    // Default rendering for circle, hollow circle.
-    if (low_quality) {
-      ctx.rect(pos[0] - 4, pos[1] - 4, 8, 8)
-    } else {
-      let radius: number
-      if (slot_shape === SlotShape.HollowCircle) {
-        doFill = false
-        doStroke = true
-        ctx.lineWidth = 3
-        ctx.strokeStyle = ctx.fillStyle
-        radius = highlight ? 4 : 3
-      } else {
-        // Normal circle
-        radius = highlight ? 5 : 4
-      }
-      ctx.arc(pos[0], pos[1], radius, 0, Math.PI * 2)
-    }
-  }
-
-  if (doFill) ctx.fill()
-  if (!low_quality && doStroke) ctx.stroke()
-
-  // render slot label
-  if (render_text) {
-    const text = slot.label || slot.localized_name || slot.name
-    if (text) {
-      // TODO: Finish impl.  Highlight text on mouseover unless we're connecting links.
-      ctx.fillStyle = label_color
-
-      if (label_position === LabelPosition.Right) {
-        if (horizontal || slot.dir == LinkDirection.UP) {
-          ctx.fillText(text, pos[0], pos[1] - 10)
-        } else {
-          ctx.fillText(text, pos[0] + 10, pos[1] + 5)
-        }
-      } else {
-        if (horizontal || slot.dir == LinkDirection.DOWN) {
-          ctx.fillText(text, pos[0], pos[1] - 8)
-        } else {
-          ctx.fillText(text, pos[0] - 10, pos[1] + 5)
-        }
-      }
-    }
-  }
-
-  // Restore the original fillStyle and strokeStyle
-  ctx.fillStyle = originalFillStyle
-  ctx.strokeStyle = originalStrokeStyle
-  ctx.lineWidth = originalLineWidth
-}
-
 interface IDrawSelectionBoundingOptions {
   /** The shape to render */
   shape?: RenderShape