Type DOMWidget and DOMWidgetOptions (#1517)

* Type DOMWidget and DOMWidgetOptions

* Annotate widget value type

src/extensions/core/uploadAudio.ts

@@ -5,7 +5,6 @@ import type { IWidget } from '@comfyorg/litegraph'
 import type { DOMWidget } from '@/scripts/domWidget'
 import { ComfyNodeDef } from '@/types/apiTypes'
 import { useToastStore } from '@/stores/toastStore'
-import { Widgets } from '@/types/comfy'
 
 type FolderType = 'input' | 'output' | 'temp'
 
@@ -37,7 +36,7 @@ function getResourceURL(
 
 async function uploadFile(
   audioWidget: IWidget,
-  audioUIWidget: DOMWidget<HTMLAudioElement>,
+  audioUIWidget: DOMWidget<HTMLAudioElement, string>,
   file: File,
   updateNode: boolean,
   pasted: boolean = false
@@ -95,12 +94,10 @@ app.registerExtension({
         audio.classList.add('comfy-audio')
         audio.setAttribute('name', 'media')
 
-        const audioUIWidget: DOMWidget<HTMLAudioElement> = node.addDOMWidget(
-          inputName,
-          /* name=*/ 'audioUI',
-          audio,
-          { serialize: false }
-        )
+        const audioUIWidget: DOMWidget<HTMLAudioElement, string> =
+          node.addDOMWidget(inputName, /* name=*/ 'audioUI', audio, {
+            serialize: false
+          })
 
         const isOutputNode = node.constructor.nodeData.output_node
         if (isOutputNode) {
@@ -121,7 +118,7 @@ app.registerExtension({
         }
         return { widget: audioUIWidget }
       }
-    } as Widgets
+    }
   },
   onNodeOutputsUpdated(nodeOutputs: Record<number, any>) {
     for (const [nodeId, output] of Object.entries(nodeOutputs)) {
@@ -129,7 +126,7 @@ app.registerExtension({
       if ('audio' in output) {
         const audioUIWidget = node.widgets.find(
           (w) => w.name === 'audioUI'
-        ) as unknown as DOMWidget<HTMLAudioElement>
+        ) as unknown as DOMWidget<HTMLAudioElement, string>
         const audio = output.audio[0]
         audioUIWidget.element.src = api.apiURL(
           getResourceURL(audio.subfolder, audio.filename, audio.type)
@@ -156,7 +153,7 @@ app.registerExtension({
         )
         const audioUIWidget = node.widgets.find(
           (w: IWidget) => w.name === 'audioUI'
-        ) as DOMWidget<HTMLAudioElement>
+        ) as unknown as DOMWidget<HTMLAudioElement, string>
 
         const onAudioWidgetUpdate = () => {
           audioUIWidget.element.src = api.apiURL(

src/scripts/domWidget.ts

@@ -3,6 +3,10 @@ import { useSettingStore } from '@/stores/settingStore'
 import { app, ANIM_PREVIEW_WIDGET } from './app'
 import { LGraphCanvas, LGraphNode, LiteGraph } from '@comfyorg/litegraph'
 import type { Vector4 } from '@comfyorg/litegraph'
+import {
+  ICustomWidget,
+  IWidgetOptions
+} from '@comfyorg/litegraph/dist/types/widgets'
 
 const SIZE = Symbol()
 
@@ -13,15 +17,20 @@ interface Rect {
   y: number
 }
 
-export interface DOMWidget<T = HTMLElement> {
-  type: string
+export interface DOMWidget<T extends HTMLElement, V extends object | string>
+  extends ICustomWidget<T> {
+  // All unrecognized types will be treated the same way as 'custom' in litegraph internally.
+  type: 'custom'
   name: string
   computedHeight?: number
   element?: T
-  options: any
-  value?: any
+  options: DOMWidgetOptions<T, V>
+  value: V
   y?: number
-  callback?: (value: any) => void
+  callback?: (value: V) => void
+  /**
+   * Draw the widget on the canvas.
+   */
   draw?: (
     ctx: CanvasRenderingContext2D,
     node: LGraphNode,
@@ -29,9 +38,31 @@ export interface DOMWidget<T = HTMLElement> {
     y: number,
     widgetHeight: number
   ) => void
+  /**
+   * TODO(huchenlei): Investigate when is this callback fired. `onRemove` is
+   * on litegraph's IBaseWidget definition, but not called in litegraph.
+   * Currently only called in widgetInputs.ts.
+   */
   onRemove?: () => void
 }
 
+export interface DOMWidgetOptions<
+  T extends HTMLElement,
+  V extends object | string
+> extends IWidgetOptions {
+  hideOnZoom?: boolean
+  selectOn?: string[]
+  onHide?: (widget: DOMWidget<T, V>) => void
+  getValue?: () => V
+  setValue?: (value: V) => void
+  getMinHeight?: () => number
+  getMaxHeight?: () => number
+  getHeight?: () => string | number
+  onDraw?: (widget: DOMWidget<T, V>) => void
+  beforeResize?: (this: DOMWidget<T, V>, node: LGraphNode) => void
+  afterResize?: (this: DOMWidget<T, V>, node: LGraphNode) => void
+}
+
 function intersect(a: Rect, b: Rect): Vector4 | null {
   const x = Math.max(a.x, b.x)
   const num1 = Math.min(a.x + a.width, b.x + b.width)
@@ -249,12 +280,15 @@ LGraphCanvas.prototype.computeVisibleNodes = function (): LGraphNode[] {
   return visibleNodes
 }
 
-LGraphNode.prototype.addDOMWidget = function (
+LGraphNode.prototype.addDOMWidget = function <
+  T extends HTMLElement,
+  V extends object | string
+>(
   name: string,
   type: string,
-  element: HTMLElement,
-  options: Record<string, any> = {}
-): DOMWidget {
+  element: T,
+  options: DOMWidgetOptions<T, V> = {}
+): DOMWidget<T, V> {
   options = { hideOnZoom: true, selectOn: ['focus', 'click'], ...options }
 
   if (!element.parentElement) {
@@ -280,13 +314,15 @@ LGraphNode.prototype.addDOMWidget = function (
     element.title = tooltip
   }
 
-  const widget: DOMWidget = {
+  const widget: DOMWidget<T, V> = {
+    // @ts-expect-error All unrecognized types will be treated the same way as 'custom'
+    // in litegraph internally.
     type,
     name,
-    get value() {
+    get value(): V {
       return options.getValue?.() ?? undefined
     },
-    set value(v) {
+    set value(v: V) {
       options.setValue?.(v)
       widget.callback?.(widget.value)
     },
@@ -306,8 +342,11 @@ LGraphNode.prototype.addDOMWidget = function (
       const hidden =
         (!!options.hideOnZoom && scale < 0.5) ||
         widget.computedHeight <= 0 ||
+        // @ts-expect-error Used by widgetInputs.ts
         widget.type === 'converted-widget' ||
+        // @ts-expect-error Used by groupNode.ts
         widget.type === 'hidden'
+
       element.dataset.shouldHide = hidden ? 'true' : 'false'
       const isInVisibleNodes = element.dataset.isInVisibleNodes === 'true'
       const isCollapsed = element.dataset.collapsed === 'true'