Skip to content

Commit

Permalink
Fix duplicate slot node types
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@moritz89
moritz89 committed Apr 8, 2023
1 parent 9be2c84 commit 8eb7613
Show file tree
Hide file tree
Showing 2 changed files with 167 additions and 55 deletions.
80 changes: 48 additions & 32 deletions src/litegraph.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -251,52 +251,68 @@
* @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
* @method registerSlotType
* @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();
}
}
Expand Down
142 changes: 119 additions & 23 deletions src/litegraph.test.js
View file
Original file line number Diff line number Diff line change
@@ -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);
};
});

Expand All @@ -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
Expand All @@ -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")
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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")
Expand All @@ -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
Expand All @@ -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();
})
});

0 comments on commit 8eb7613

Please sign in to comment.