From 8eb76138d0bb89d801150020ad36faa761eb7da3 Mon Sep 17 00:00:00 2001 From: Moritz Ulmer Date: Sat, 8 Apr 2023 20:59:38 +0200 Subject: [PATCH] Fix duplicate slot node types Why: - If a node has multiple of the same slot types, it is duplicated This change addresses the need by: - Checking if the node type has been added for a slot type - Adding tests - Improve legibility --- src/litegraph.js | 80 ++++++++++++++---------- src/litegraph.test.js | 142 +++++++++++++++++++++++++++++++++++------- 2 files changed, 167 insertions(+), 55 deletions(-) diff --git a/src/litegraph.js b/src/litegraph.js index a70f4733e..a82c39f3c 100644 --- a/src/litegraph.js +++ b/src/litegraph.js @@ -251,13 +251,18 @@ * @param {String|Object} type name of the node or the node constructor itself */ unregisterNodeType: function(type) { - var base_class = type.constructor === String ? this.registered_node_types[type] : type; - if(!base_class) - throw("node type not found: " + type ); - delete this.registered_node_types[base_class.type]; - if(base_class.constructor.name) - delete this.Nodes[base_class.constructor.name]; - }, + const base_class = + type.constructor === String + ? this.registered_node_types[type] + : type; + if (!base_class) { + throw "node type not found: " + type; + } + delete this.registered_node_types[base_class.type]; + if (base_class.constructor.name) { + delete this.Nodes[base_class.constructor.name]; + } + }, /** * Save a slot type and his node @@ -265,38 +270,49 @@ * @param {String|Object} type name of the node or the node constructor itself * @param {String} slot_type name of the slot type (variable type), eg. string, number, array, boolean, .. */ - registerNodeAndSlotType: function(type,slot_type,out){ + registerNodeAndSlotType: function(type, slot_type, out){ out = out || false; - var base_class = type.constructor === String && this.registered_node_types[type] !== "anonymous" ? this.registered_node_types[type] : type; - - var sCN = base_class.constructor.type; - - if (typeof slot_type == "string"){ - var aTypes = slot_type.split(","); - }else if (slot_type == this.EVENT || slot_type == this.ACTION){ - var aTypes = ["_event_"]; - }else{ - var aTypes = ["*"]; + const base_class = + type.constructor === String && + this.registered_node_types[type] !== "anonymous" + ? this.registered_node_types[type] + : type; + + const class_type = base_class.constructor.type; + + let allTypes = []; + if (typeof slot_type === "string") { + allTypes = slot_type.split(","); + } else if (slot_type == this.EVENT || slot_type == this.ACTION) { + allTypes = ["_event_"]; + } else { + allTypes = ["*"]; } - for (var i = 0; i < aTypes.length; ++i) { - var sT = aTypes[i]; //.toLowerCase(); - if (sT === ""){ - sT = "*"; + for (let i = 0; i < allTypes.length; ++i) { + let slotType = allTypes[i]; + if (slotType === "") { + slotType = "*"; } - var registerTo = out ? "registered_slot_out_types" : "registered_slot_in_types"; - if (typeof this[registerTo][sT] == "undefined") this[registerTo][sT] = {nodes: []}; - this[registerTo][sT].nodes.push(sCN); - + const registerTo = out + ? "registered_slot_out_types" + : "registered_slot_in_types"; + if (this[registerTo][slotType] === undefined) { + this[registerTo][slotType] = { nodes: [] }; + } + if (!this[registerTo][slotType].nodes.includes(class_type)) { + this[registerTo][slotType].nodes.push(class_type); + } + // check if is a new type - if (!out){ - if (!this.slot_types_in.includes(sT.toLowerCase())){ - this.slot_types_in.push(sT.toLowerCase()); + if (!out) { + if (!this.slot_types_in.includes(slotType.toLowerCase())) { + this.slot_types_in.push(slotType.toLowerCase()); this.slot_types_in.sort(); } - }else{ - if (!this.slot_types_out.includes(sT.toLowerCase())){ - this.slot_types_out.push(sT.toLowerCase()); + } else { + if (!this.slot_types_out.includes(slotType.toLowerCase())) { + this.slot_types_out.push(slotType.toLowerCase()); this.slot_types_out.sort(); } } diff --git a/src/litegraph.test.js b/src/litegraph.test.js index 384f37bb1..a5c034205 100644 --- a/src/litegraph.test.js +++ b/src/litegraph.test.js @@ -1,12 +1,17 @@ describe("register node types", () => { let lg; - let calc_sum; + let Sum; beforeEach(() => { jest.resetModules(); lg = require("./litegraph"); - calc_sum = function calc_sum(a, b) { - return a + b; + Sum = function Sum() { + this.addInput("a", "number"); + this.addInput("b", "number"); + this.addOutput("sum", "number"); + }; + Sum.prototype.onExecute = function (a, b) { + this.setOutputData(0, a + b); }; }); @@ -15,12 +20,12 @@ describe("register node types", () => { }); test("normal case", () => { - lg.LiteGraph.registerNodeType("math/sum", calc_sum); + lg.LiteGraph.registerNodeType("math/sum", Sum); let node = lg.LiteGraph.registered_node_types["math/sum"]; expect(node).toBeTruthy(); expect(node.type).toBe("math/sum"); - expect(node.title).toBe("calc_sum"); + expect(node.title).toBe("Sum"); expect(node.category).toBe("math"); expect(node.prototype.configure).toBe( lg.LGraphNode.prototype.configure @@ -34,10 +39,10 @@ describe("register node types", () => { lg.LiteGraph.onNodeTypeRegistered = jest.fn(); lg.LiteGraph.onNodeTypeReplaced = jest.fn(); - lg.LiteGraph.registerNodeType("math/sum", calc_sum); + lg.LiteGraph.registerNodeType("math/sum", Sum); expect(lg.LiteGraph.onNodeTypeRegistered).toHaveBeenCalled(); expect(lg.LiteGraph.onNodeTypeReplaced).not.toHaveBeenCalled(); - lg.LiteGraph.registerNodeType("math/sum", calc_sum); + lg.LiteGraph.registerNodeType("math/sum", Sum); expect(lg.LiteGraph.onNodeTypeReplaced).toHaveBeenCalled(); expect(consoleSpy).toHaveBeenCalledWith( expect.stringMatching("replacing node type") @@ -48,8 +53,8 @@ describe("register node types", () => { }); test("node with title", () => { - calc_sum.title = "The sum title"; - lg.LiteGraph.registerNodeType("math/sum", calc_sum); + Sum.title = "The sum title"; + lg.LiteGraph.registerNodeType("math/sum", Sum); let node = lg.LiteGraph.registered_node_types["math/sum"]; expect(node.title).toBe("The sum title"); expect(node.title).not.toBe(node.name); @@ -62,7 +67,7 @@ describe("register node types", () => { }); test("check shape mapping", () => { - lg.LiteGraph.registerNodeType("math/sum", calc_sum); + lg.LiteGraph.registerNodeType("math/sum", Sum); const node_type = lg.LiteGraph.registered_node_types["math/sum"]; expect(new node_type().shape).toBe(undefined); @@ -81,10 +86,10 @@ describe("register node types", () => { // Check that also works for replaced node types jest.spyOn(console, "log").mockImplementation(() => {}); - function new_calc_sum(a, b) { + function NewCalcSum(a, b) { return a + b; } - lg.LiteGraph.registerNodeType("math/sum", new_calc_sum); + lg.LiteGraph.registerNodeType("math/sum", NewCalcSum); const new_node_type = lg.LiteGraph.registered_node_types["math/sum"]; new_node_type.prototype.shape = "box"; expect(new new_node_type().shape).toBe(lg.LiteGraph.BOX_SHAPE); @@ -95,8 +100,8 @@ describe("register node types", () => { .spyOn(console, "warn") .mockImplementation(() => {}); - calc_sum.prototype.onPropertyChange = true; - lg.LiteGraph.registerNodeType("math/sum", calc_sum); + Sum.prototype.onPropertyChange = true; + lg.LiteGraph.registerNodeType("math/sum", Sum); expect(consoleSpy).toBeCalledTimes(1); expect(consoleSpy).toBeCalledWith( expect.stringContaining("has onPropertyChange method") @@ -108,13 +113,18 @@ describe("register node types", () => { expect(lg.LiteGraph.node_types_by_file_extension).toEqual({}); // Create two node types with calc_times overriding .pdf - calc_sum.supported_extensions = ["PDF", "exe", null]; - function calc_times(a, b) { - return a * b; + Sum.supported_extensions = ["PDF", "exe", null]; + function Times() { + this.addInput("a", "number"); + this.addInput("b", "number"); + this.addOutput("times", "number"); } - calc_times.supported_extensions = ["pdf", "jpg"]; - lg.LiteGraph.registerNodeType("math/sum", calc_sum); - lg.LiteGraph.registerNodeType("math/times", calc_times); + Times.prototype.onExecute = function (a, b) { + this.setOutputData(0, a * b); + }; + Times.supported_extensions = ["pdf", "jpg"]; + lg.LiteGraph.registerNodeType("math/sum", Sum); + lg.LiteGraph.registerNodeType("math/times", Times); expect( Object.keys(lg.LiteGraph.node_types_by_file_extension).length @@ -123,8 +133,94 @@ describe("register node types", () => { expect(lg.LiteGraph.node_types_by_file_extension).toHaveProperty("exe"); expect(lg.LiteGraph.node_types_by_file_extension).toHaveProperty("jpg"); - expect(lg.LiteGraph.node_types_by_file_extension.exe).toBe(calc_sum); - expect(lg.LiteGraph.node_types_by_file_extension.pdf).toBe(calc_times); - expect(lg.LiteGraph.node_types_by_file_extension.jpg).toBe(calc_times); + expect(lg.LiteGraph.node_types_by_file_extension.exe).toBe(Sum); + expect(lg.LiteGraph.node_types_by_file_extension.pdf).toBe(Times); + expect(lg.LiteGraph.node_types_by_file_extension.jpg).toBe(Times); + }); + + test("register in/out slot types", () => { + expect(lg.LiteGraph.registered_slot_in_types).toEqual({}); + expect(lg.LiteGraph.registered_slot_out_types).toEqual({}); + + // Test slot type registration with first type + lg.LiteGraph.auto_load_slot_types = true; + lg.LiteGraph.registerNodeType("math/sum", Sum); + expect(lg.LiteGraph.registered_slot_in_types).toEqual({ + number: { nodes: ["math/sum"] }, + }); + expect(lg.LiteGraph.registered_slot_out_types).toEqual({ + number: { nodes: ["math/sum"] }, + }); + + // Test slot type registration with second type + function ToInt() { + this.addInput("string", "string"); + this.addOutput("number", "number"); + }; + ToInt.prototype.onExecute = function (str) { + this.setOutputData(0, Number(str)); + }; + lg.LiteGraph.registerNodeType("basic/to_int", ToInt); + expect(lg.LiteGraph.registered_slot_in_types).toEqual({ + number: { nodes: ["math/sum"] }, + string: { nodes: ["basic/to_int"] }, + }); + expect(lg.LiteGraph.registered_slot_out_types).toEqual({ + number: { nodes: ["math/sum", "basic/to_int"] }, + }); + }); +}); + +describe("unregister node types", () => { + let lg; + let Sum; + + beforeEach(() => { + jest.resetModules(); + lg = require("./litegraph"); + Sum = function Sum() { + this.addInput("a", "number"); + this.addInput("b", "number"); + this.addOutput("sum", "number"); + }; + Sum.prototype.onExecute = function (a, b) { + this.setOutputData(0, a + b); + }; + }); + + afterEach(() => { + jest.restoreAllMocks(); + }); + + test("remove by name", () => { + lg.LiteGraph.registerNodeType("math/sum", Sum); + expect(lg.LiteGraph.registered_node_types["math/sum"]).toBeTruthy(); + + lg.LiteGraph.unregisterNodeType("math/sum"); + expect(lg.LiteGraph.registered_node_types["math/sum"]).toBeFalsy(); }); + + test("remove by object", () => { + lg.LiteGraph.registerNodeType("math/sum", Sum); + expect(lg.LiteGraph.registered_node_types["math/sum"]).toBeTruthy(); + + lg.LiteGraph.unregisterNodeType(Sum); + expect(lg.LiteGraph.registered_node_types["math/sum"]).toBeFalsy(); + }); + + test("try removing with wrong name", () => { + expect(() => lg.LiteGraph.unregisterNodeType("missing/type")).toThrow( + "node type not found: missing/type" + ); + }); + + test("no constructor name", () => { + function BlankNode() {} + BlankNode.constructor = {} + lg.LiteGraph.registerNodeType("blank/node", BlankNode); + expect(lg.LiteGraph.registered_node_types["blank/node"]).toBeTruthy() + + lg.LiteGraph.unregisterNodeType("blank/node"); + expect(lg.LiteGraph.registered_node_types["blank/node"]).toBeFalsy(); + }) });