From c75157e86e832c8eeeea0ae6bf9066a767512a71 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Sun, 9 Feb 2025 15:23:38 -0500
Subject: [PATCH] Convert i/o slot to class in LGraphNode.configure (#506)

---
 src/LGraphNode.ts       | 50 +++++++++++++++++++----------------------
 test/LGraphNode.test.ts | 36 +++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 27 deletions(-)

diff --git a/src/LGraphNode.ts b/src/LGraphNode.ts
index b1f7e6403..1908bdbf7 100644
--- a/src/LGraphNode.ts
+++ b/src/LGraphNode.ts
@@ -596,44 +596,40 @@ export class LGraphNode implements Positionable, IPinnable {
       this.title = this.constructor.title
     }
 
-    if (this.inputs) {
-      for (let i = 0; i < this.inputs.length; ++i) {
-        const input = this.inputs[i]
-        const link = this.graph ? this.graph._links.get(input.link) : null
-        this.onConnectionsChange?.(NodeSlotType.INPUT, i, true, link, input)
-        this.onInputAdded?.(input)
-      }
+    this.inputs ??= []
+    this.inputs = this.inputs.map(input => toClass(NodeInputSlot, input))
+    for (const [i, input] of this.inputs.entries()) {
+      const link = this.graph ? this.graph._links.get(input.link) : null
+      this.onConnectionsChange?.(NodeSlotType.INPUT, i, true, link, input)
+      this.onInputAdded?.(input)
     }
 
-    if (this.outputs) {
-      for (let i = 0; i < this.outputs.length; ++i) {
-        const output = this.outputs[i]
-        if (!output.links) {
-          continue
-        }
-        for (let j = 0; j < output.links.length; ++j) {
-          const link = this.graph
-            ? this.graph._links.get(output.links[j])
-            : null
-          this.onConnectionsChange?.(NodeSlotType.OUTPUT, i, true, link, output)
-        }
-        this.onOutputAdded?.(output)
+    this.outputs ??= []
+    this.outputs = this.outputs.map(output => toClass(NodeOutputSlot, output))
+    for (const [i, output] of this.outputs.entries()) {
+      if (!output.links) {
+        continue
       }
+      for (const linkId of output.links) {
+        const link = this.graph
+          ? this.graph._links.get(linkId)
+          : null
+        this.onConnectionsChange?.(NodeSlotType.OUTPUT, i, true, link, output)
+      }
+      this.onOutputAdded?.(output)
     }
 
     if (this.widgets) {
-      for (let i = 0; i < this.widgets.length; ++i) {
-        const w = this.widgets[i]
+      for (const w of this.widgets) {
         if (!w) continue
 
         if (w.options?.property && this.properties[w.options.property] != undefined)
           w.value = JSON.parse(JSON.stringify(this.properties[w.options.property]))
       }
-      if (info.widgets_values) {
-        for (let i = 0; i < info.widgets_values.length; ++i) {
-          if (this.widgets[i]) {
-            this.widgets[i].value = info.widgets_values[i]
-          }
+
+      for (const [i, value] of (info.widgets_values ?? []).entries()) {
+        if (this.widgets[i]) {
+          this.widgets[i].value = value
         }
       }
     }
diff --git a/test/LGraphNode.test.ts b/test/LGraphNode.test.ts
index f1634a09f..3bc249b6a 100644
--- a/test/LGraphNode.test.ts
+++ b/test/LGraphNode.test.ts
@@ -1,5 +1,6 @@
 import { describe, expect } from "vitest"
 import { LGraphNode } from "@/litegraph"
+import { NodeInputSlot, NodeOutputSlot } from "@/NodeSlot"
 import { lgTest } from "./lgTest"
 
 describe("LGraphNode", () => {
@@ -13,4 +14,39 @@ describe("LGraphNode", () => {
     expect(node.size).toEqual(new Float32Array([100, 100]))
     expect(node.serialize().size).toEqual([100, 100])
   })
+
+  lgTest("should configure inputs correctly", () => {
+    const node = new LGraphNode("TestNode")
+    node.configure({
+      id: 0,
+      inputs: [{ name: "TestInput", type: "number", link: null }],
+    })
+    expect(node.inputs.length).toEqual(1)
+    expect(node.inputs[0].name).toEqual("TestInput")
+    expect(node.inputs[0].link).toEqual(null)
+    expect(node.inputs[0]).instanceOf(NodeInputSlot)
+
+    // Should not override existing inputs
+    node.configure({ id: 1 })
+    expect(node.id).toEqual(1)
+    expect(node.inputs.length).toEqual(1)
+  })
+
+  lgTest("should configure outputs correctly", () => {
+    const node = new LGraphNode("TestNode")
+    node.configure({
+      id: 0,
+      outputs: [{ name: "TestOutput", type: "number", links: [] }],
+    })
+    expect(node.outputs.length).toEqual(1)
+    expect(node.outputs[0].name).toEqual("TestOutput")
+    expect(node.outputs[0].type).toEqual("number")
+    expect(node.outputs[0].links).toEqual([])
+    expect(node.outputs[0]).instanceOf(NodeOutputSlot)
+
+    // Should not override existing outputs
+    node.configure({ id: 1 })
+    expect(node.id).toEqual(1)
+    expect(node.outputs.length).toEqual(1)
+  })
 })