From 608b5f83426bd510e0dddbe0b828193e79e0e5a3 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Fri, 7 Feb 2025 17:13:32 -0500 Subject: [PATCH] Implement NodeSlot (#476) * Implement NodeSlot * nit --- src/LGraphCanvas.ts | 30 +++----------- src/LGraphNode.ts | 9 +++-- src/NodeSlot.ts | 95 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 106 insertions(+), 28 deletions(-) create mode 100644 src/NodeSlot.ts diff --git a/src/LGraphCanvas.ts b/src/LGraphCanvas.ts index 96b8293c7b..0105bd864e 100644 --- a/src/LGraphCanvas.ts +++ b/src/LGraphCanvas.ts @@ -67,6 +67,7 @@ import { getAllNestedItems, findFirstNode } from "./utils/collections" import { CanvasPointer } from "./CanvasPointer" import { BooleanWidget } from "./widgets/BooleanWidget" import { toClass } from "./utils/type" +import { NodeInputSlot, NodeOutputSlot, type ConnectionColorContext } from "./NodeSlot" interface IShowSearchOptions { node_to?: LGraphNode @@ -177,7 +178,7 @@ interface IPasteFromClipboardOptions { * This class is in charge of rendering one graph inside a canvas. And provides all the interaction required. * Valid callbacks are: onNodeSelected, onNodeDeselected, onShowNodePanel, onNodeDblClicked */ -export class LGraphCanvas { +export class LGraphCanvas implements ConnectionColorContext { // Optimised buffers used during rendering static #temp = new Float32Array(4) static #temp_vec2 = new Float32Array(2) @@ -4877,9 +4878,7 @@ export class LGraphCanvas { // input connection slots if (node.inputs) { for (let i = 0; i < node.inputs.length; i++) { - const slot = node.inputs[i] - - const slot_type = slot.type + const slot = toClass(NodeInputSlot, node.inputs[i]) // change opacity of incompatible slots when dragging a connection const isValid = @@ -4891,15 +4890,7 @@ export class LGraphCanvas { : LiteGraph.NODE_TEXT_COLOR ctx.globalAlpha = isValid ? editor_alpha : 0.4 * editor_alpha - ctx.fillStyle = - slot.link != null - ? slot.color_on || - this.default_connection_color_byType[slot_type] || - this.default_connection_color.input_on - : slot.color_off || - this.default_connection_color_byTypeOff[slot_type] || - this.default_connection_color_byType[slot_type] || - this.default_connection_color.input_off + ctx.fillStyle = slot.renderingColor(this) const pos = node.getConnectionPos(true, i, slot_pos) pos[0] -= node.pos[0] @@ -4926,7 +4917,7 @@ export class LGraphCanvas { ctx.strokeStyle = "black" if (node.outputs) { for (let i = 0; i < node.outputs.length; i++) { - const slot = node.outputs[i] + const slot = toClass(NodeOutputSlot, node.outputs[i]) const slot_type = slot.type @@ -4947,16 +4938,7 @@ export class LGraphCanvas { max_y = pos[1] + LiteGraph.NODE_SLOT_HEIGHT * 0.5 } - ctx.fillStyle = - slot.links && slot.links.length - ? slot.color_on || - this.default_connection_color_byType[slot_type] || - this.default_connection_color.output_on - : slot.color_off || - this.default_connection_color_byTypeOff[slot_type] || - this.default_connection_color_byType[slot_type] || - this.default_connection_color.output_off - + ctx.fillStyle = slot.renderingColor(this) drawSlot(ctx, slot, pos, { horizontal, low_quality, diff --git a/src/LGraphNode.ts b/src/LGraphNode.ts index b4fcd4694e..8d9489cdca 100644 --- a/src/LGraphNode.ts +++ b/src/LGraphNode.ts @@ -33,6 +33,7 @@ import { type LGraphNodeConstructor, LiteGraph } from "./litegraph" import { isInRectangle, isInRect, snapPoint } from "./measure" import { LLink } from "./LLink" import { BooleanWidget } from "./widgets/BooleanWidget" +import { NodeInputSlot, NodeOutputSlot } from "./NodeSlot" export type NodeId = number | string @@ -1308,7 +1309,7 @@ export class LGraphNode implements Positionable, IPinnable { type?: ISlotType, extra_info?: object, ): INodeOutputSlot { - const output = { name: name, type: type, links: null } + const output = new NodeOutputSlot({ name: name, type: type, links: null }) if (extra_info) { for (const i in extra_info) { output[i] = extra_info[i] @@ -1334,7 +1335,7 @@ export class LGraphNode implements Positionable, IPinnable { addOutputs(array: [string, ISlotType, Record][]): void { for (let i = 0; i < array.length; ++i) { const info = array[i] - const o = { name: info[0], type: info[1], links: null } + const o = new NodeOutputSlot({ name: info[0], type: info[1], links: null }) if (array[2]) { for (const j in info[2]) { o[j] = info[2][j] @@ -1383,7 +1384,7 @@ export class LGraphNode implements Positionable, IPinnable { */ addInput(name: string, type: ISlotType, extra_info?: object): INodeInputSlot { type = type || 0 - const input: INodeInputSlot = { name: name, type: type, link: null } + const input: INodeInputSlot = new NodeInputSlot({ name: name, type: type, link: null }) if (extra_info) { for (const i in extra_info) { input[i] = extra_info[i] @@ -1408,7 +1409,7 @@ export class LGraphNode implements Positionable, IPinnable { addInputs(array: [string, ISlotType, Record][]): void { for (let i = 0; i < array.length; ++i) { const info = array[i] - const o: INodeInputSlot = { name: info[0], type: info[1], link: null } + const o: INodeInputSlot = new NodeInputSlot({ name: info[0], type: info[1], link: null }) // TODO: Checking the wrong variable here - confirm no downstream consumers, then remove. if (array[2]) { for (const j in info[2]) { diff --git a/src/NodeSlot.ts b/src/NodeSlot.ts new file mode 100644 index 0000000000..433f565cda --- /dev/null +++ b/src/NodeSlot.ts @@ -0,0 +1,95 @@ +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" + +export interface ConnectionColorContext { + default_connection_color: { + input_off: string + input_on: string + output_off: string + output_on: string + } + default_connection_color_byType: Dictionary + default_connection_color_byTypeOff: Dictionary +} + +export abstract class NodeSlot implements INodeSlot { + name: string + localized_name?: string + label?: string + type: ISlotType + dir?: LinkDirection + removable?: boolean + shape?: RenderShape + color_off?: CanvasColour + color_on?: CanvasColour + locked?: boolean + nameLocked?: boolean + pos?: Point + widget?: IWidget + + constructor(slot: INodeSlot) { + Object.assign(this, slot) + this.name = slot.name + this.type = slot.type + } + + /** + * The label to display in the UI. + */ + get displayLabel(): string { + return this.label || this.localized_name || this.name || "" + } + + abstract isConnected(): boolean + + connectedColor(context: ConnectionColorContext): CanvasColour { + return this.color_on || + context.default_connection_color_byType[this.type] || + context.default_connection_color.output_on + } + + disconnectedColor(context: ConnectionColorContext): CanvasColour { + return this.color_off || + context.default_connection_color_byTypeOff[this.type] || + context.default_connection_color_byType[this.type] || + context.default_connection_color.output_off + } + + renderingColor(context: ConnectionColorContext): CanvasColour { + return this.isConnected() + ? this.connectedColor(context) + : this.disconnectedColor(context) + } +} + +export class NodeInputSlot extends NodeSlot implements INodeInputSlot { + link: LinkId | null + + constructor(slot: INodeInputSlot) { + super(slot) + this.link = slot.link + } + + override isConnected(): boolean { + return this.link != null + } +} + +export class NodeOutputSlot extends NodeSlot implements INodeOutputSlot { + links: LinkId[] | null + _data?: unknown + slot_index?: number + + constructor(slot: INodeOutputSlot) { + super(slot) + this.links = slot.links + this._data = slot._data + this.slot_index = slot.slot_index + } + + override isConnected(): boolean { + return this.links != null && this.links.length > 0 + } +}