From 01b8ae5bb8ee5a15f00612b678261a397b94b9ab Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Sat, 8 Feb 2025 23:41:01 -0500
Subject: [PATCH] [Refactor] LGraphNode.drawWidgets (#499)

---
 src/LGraphCanvas.ts | 65 ++++++++-------------------------------------
 src/LGraphNode.ts   | 62 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 55 deletions(-)

diff --git a/src/LGraphCanvas.ts b/src/LGraphCanvas.ts
index 31eb169dd..a4f0fc461 100644
--- a/src/LGraphCanvas.ts
+++ b/src/LGraphCanvas.ts
@@ -4780,18 +4780,7 @@ export class LGraphCanvas implements ConnectionColorContext {
       ctx.textAlign = "left"
       ctx.globalAlpha = 1
 
-      if (node.widgets) {
-        let widgets_y = max_y
-        if (horizontal || node.widgets_up) {
-          widgets_y = 2
-        }
-        if (node.widgets_start_y != null) widgets_y = node.widgets_start_y
-        this.drawNodeWidgets(
-          node,
-          widgets_y,
-          ctx,
-        )
-      }
+      this.drawNodeWidgets(node, max_y, ctx)
     } else if (this.render_collapsed_slots) {
       // if collapsed
       let input_slot = null
@@ -5679,54 +5668,22 @@ export class LGraphCanvas implements ConnectionColorContext {
 
   /**
    * draws the widgets stored inside a node
+   * @deprecated Use {@link LGraphNode.drawWidgets} instead.
+   * @note Currently there are extensions hijacking this function, so we cannot remove it.
    */
   drawNodeWidgets(
     node: LGraphNode,
     posY: number,
     ctx: CanvasRenderingContext2D,
   ): void {
-    if (!node.widgets || !node.widgets.length) return
-    const width = node.size[0]
-    const widgets = node.widgets
-    posY += 2
-    const H = LiteGraph.NODE_WIDGET_HEIGHT
-    const show_text = !this.low_quality
-    ctx.save()
-    ctx.globalAlpha = this.editor_alpha
-    const margin = 15
-
-    for (const w of widgets) {
-      if (w.hidden || (w.advanced && !node.showAdvanced)) continue
-      const y = w.y || posY
-      const outline_color = w.advanced ? LiteGraph.WIDGET_ADVANCED_OUTLINE_COLOR : LiteGraph.WIDGET_OUTLINE_COLOR
-
-      if (w === this.link_over_widget) {
-        // Manually draw a slot next to the widget simulating an input
-        new NodeInputSlot({
-          name: "",
-          type: this.link_over_widget_type,
-          link: 0,
-        }).draw(ctx, { pos: [10, y + 10], colorContext: this })
-      }
-
-      w.last_y = y
-      ctx.strokeStyle = outline_color
-      ctx.fillStyle = "#222"
-      ctx.textAlign = "left"
-      if (w.disabled) ctx.globalAlpha *= 0.5
-      const widget_width = w.width || width
-
-      const WidgetClass = WIDGET_TYPE_MAP[w.type]
-      if (WidgetClass) {
-        toClass(WidgetClass, w).drawWidget(ctx, { y, width: widget_width, show_text, margin })
-      } else {
-        w.draw?.(ctx, node, widget_width, y, H)
-      }
-      posY += (w.computeSize ? w.computeSize(widget_width)[1] : H) + 4
-      ctx.globalAlpha = this.editor_alpha
-    }
-    ctx.restore()
-    ctx.textAlign = "left"
+    node.drawWidgets(ctx, {
+      y: posY,
+      colorContext: this,
+      linkOverWidget: this.link_over_widget,
+      linkOverWidgetType: this.link_over_widget_type,
+      lowQuality: this.low_quality,
+      editorAlpha: this.editor_alpha,
+    })
   }
 
   /**
diff --git a/src/LGraphNode.ts b/src/LGraphNode.ts
index a176820ff..e132a9682 100644
--- a/src/LGraphNode.ts
+++ b/src/LGraphNode.ts
@@ -32,8 +32,9 @@ import { BadgePosition, LGraphBadge } from "./LGraphBadge"
 import { type LGraphNodeConstructor, LiteGraph } from "./litegraph"
 import { isInRectangle, isInRect, snapPoint } from "./measure"
 import { LLink } from "./LLink"
-import { NodeInputSlot, NodeOutputSlot } from "./NodeSlot"
+import { ConnectionColorContext, NodeInputSlot, NodeOutputSlot } from "./NodeSlot"
 import { WIDGET_TYPE_MAP } from "./widgets/widgetMap"
+import { toClass } from "./utils/type"
 export type NodeId = number | string
 
 export interface INodePropertyInfo {
@@ -3086,4 +3087,63 @@ export class LGraphNode implements Positionable, IPinnable {
       }
     }
   }
+
+  drawWidgets(ctx: CanvasRenderingContext2D, options: {
+    y: number
+    colorContext: ConnectionColorContext
+    linkOverWidget: IWidget
+    linkOverWidgetType: ISlotType
+    lowQuality?: boolean
+    editorAlpha?: number
+  }): void {
+    if (!this.widgets) return
+
+    const { y, colorContext, linkOverWidget, linkOverWidgetType, lowQuality = false, editorAlpha = 1 } = options
+    let posY = y
+    if (this.horizontal || this.widgets_up) {
+      posY = 2
+    }
+    if (this.widgets_start_y != null) posY = this.widgets_start_y
+
+    const width = this.size[0]
+    const widgets = this.widgets
+    posY += 2
+    const H = LiteGraph.NODE_WIDGET_HEIGHT
+    const show_text = !lowQuality
+    ctx.save()
+    ctx.globalAlpha = editorAlpha
+    const margin = 15
+
+    for (const w of widgets) {
+      if (w.hidden || (w.advanced && !this.showAdvanced)) continue
+      const y = w.y || posY
+      const outline_color = w.advanced ? LiteGraph.WIDGET_ADVANCED_OUTLINE_COLOR : LiteGraph.WIDGET_OUTLINE_COLOR
+
+      if (w === linkOverWidget) {
+        // Manually draw a slot next to the widget simulating an input
+        new NodeInputSlot({
+          name: "",
+          type: linkOverWidgetType,
+          link: 0,
+        }).draw(ctx, { pos: [10, y + 10], colorContext })
+      }
+
+      w.last_y = y
+      ctx.strokeStyle = outline_color
+      ctx.fillStyle = "#222"
+      ctx.textAlign = "left"
+      if (w.disabled) ctx.globalAlpha *= 0.5
+      const widget_width = w.width || width
+
+      const WidgetClass = WIDGET_TYPE_MAP[w.type]
+      if (WidgetClass) {
+        toClass(WidgetClass, w).drawWidget(ctx, { y, width: widget_width, show_text, margin })
+      } else {
+        w.draw?.(ctx, this, widget_width, y, H)
+      }
+      posY += (w.computeSize ? w.computeSize(widget_width)[1] : H) + 4
+      ctx.globalAlpha = editorAlpha
+    }
+    ctx.restore()
+  }
 }