From 8e293b41f565c4f4dc1f6483f5818486172fcf3b Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Tue, 3 Dec 2024 08:39:29 -0800
Subject: [PATCH] [Refactor] Remove abstract inheritance of node search filter
 (#1774)

---
 .../searchbox/NodeSearchBoxPopover.vue        |   2 +-
 src/services/nodeSearchService.ts             | 124 +++++++++---------
 2 files changed, 63 insertions(+), 63 deletions(-)

diff --git a/src/components/searchbox/NodeSearchBoxPopover.vue b/src/components/searchbox/NodeSearchBoxPopover.vue
index aa58098ba..f16a318fe 100644
--- a/src/components/searchbox/NodeSearchBoxPopover.vue
+++ b/src/components/searchbox/NodeSearchBoxPopover.vue
@@ -22,7 +22,7 @@
     >
       <template #container>
         <NodeSearchBox
-          :filters="nodeFilters as FilterAndValue[]"
+          :filters="nodeFilters"
           @add-filter="addFilter"
           @remove-filter="removeFilter"
           @add-node="addNode"
diff --git a/src/services/nodeSearchService.ts b/src/services/nodeSearchService.ts
index c57e77217..a1c60ff20 100644
--- a/src/services/nodeSearchService.ts
+++ b/src/services/nodeSearchService.ts
@@ -10,8 +10,8 @@ interface ExtraSearchOptions {
 }
 
 export class FuseSearch<T> {
-  private fuse: Fuse<T>
-  private readonly keys: string[]
+  public readonly fuse: Fuse<T>
+  public readonly keys: string[]
   public readonly data: T[]
   public readonly advancedScoring: boolean
 
@@ -130,21 +130,30 @@ export class FuseSearch<T> {
 
 export type FilterAndValue<T = string> = [NodeFilter<T>, T]
 
-export abstract class NodeFilter<FilterOptionT = string> {
-  public abstract readonly id: string
-  public abstract readonly name: string
-  public abstract readonly invokeSequence: string
-  public abstract readonly longInvokeSequence: string
+export class NodeFilter<FilterOptionT = string> {
   public readonly fuseSearch: FuseSearch<FilterOptionT>
 
   constructor(
+    public readonly id: string,
+    public readonly name: string,
+    public readonly invokeSequence: string,
+    public readonly longInvokeSequence: string,
+    public readonly nodeOptions:
+      | FilterOptionT[]
+      | ((node: ComfyNodeDefImpl) => FilterOptionT[]),
     nodeDefs: ComfyNodeDefImpl[],
     options?: IFuseOptions<FilterOptionT>
   ) {
     this.fuseSearch = new FuseSearch(this.getAllNodeOptions(nodeDefs), options)
   }
 
-  private getAllNodeOptions(nodeDefs: ComfyNodeDefImpl[]): FilterOptionT[] {
+  public getNodeOptions(node: ComfyNodeDefImpl): FilterOptionT[] {
+    return this.nodeOptions instanceof Function
+      ? this.nodeOptions(node)
+      : this.nodeOptions
+  }
+
+  public getAllNodeOptions(nodeDefs: ComfyNodeDefImpl[]): FilterOptionT[] {
     return [
       ...new Set(
         nodeDefs.reduce((acc, nodeDef) => {
@@ -154,8 +163,6 @@ export abstract class NodeFilter<FilterOptionT = string> {
     ]
   }
 
-  public abstract getNodeOptions(node: ComfyNodeDefImpl): FilterOptionT[]
-
   public matches(
     node: ComfyNodeDefImpl,
     value: FilterOptionT,
@@ -173,50 +180,6 @@ export abstract class NodeFilter<FilterOptionT = string> {
   }
 }
 
-export class InputTypeFilter extends NodeFilter<string> {
-  public readonly id: string = 'input'
-  public readonly name = 'Input Type'
-  public readonly invokeSequence = 'i'
-  public readonly longInvokeSequence = 'input'
-
-  public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
-    return node.input.all.map((input) => input.type)
-  }
-}
-
-export class OutputTypeFilter extends NodeFilter<string> {
-  public readonly id: string = 'output'
-  public readonly name = 'Output Type'
-  public readonly invokeSequence = 'o'
-  public readonly longInvokeSequence = 'output'
-
-  public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
-    return node.output.all.map((output) => output.type)
-  }
-}
-
-export class NodeSourceFilter extends NodeFilter<string> {
-  public readonly id: string = 'source'
-  public readonly name = 'Source'
-  public readonly invokeSequence = 's'
-  public readonly longInvokeSequence = 'source'
-
-  public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
-    return [node.nodeSource.displayText]
-  }
-}
-
-export class NodeCategoryFilter extends NodeFilter<string> {
-  public readonly id: string = 'category'
-  public readonly name = 'Category'
-  public readonly invokeSequence = 'c'
-  public readonly longInvokeSequence = 'category'
-
-  public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
-    return [node.category]
-  }
-}
-
 export class NodeSearchService {
   public readonly nodeFuseSearch: FuseSearch<ComfyNodeDefImpl>
   public readonly nodeFilters: NodeFilter<string>[]
@@ -241,15 +204,52 @@ export class NodeSearchService {
       shouldSort: true
     }
 
-    this.nodeFilters = [
-      new InputTypeFilter(data, filterSearchOptions),
-      new OutputTypeFilter(data, filterSearchOptions),
-      new NodeCategoryFilter(data, filterSearchOptions)
-    ]
+    const inputTypeFilter = new NodeFilter<string>(
+      /* id */ 'input',
+      /* name */ 'Input Type',
+      /* invokeSequence */ 'i',
+      /* longInvokeSequence */ 'input',
+      (node) => node.input.all.map((input) => input.type),
+      data,
+      filterSearchOptions
+    )
 
-    if (data[0].python_module !== undefined) {
-      this.nodeFilters.push(new NodeSourceFilter(data, filterSearchOptions))
-    }
+    const outputTypeFilter = new NodeFilter<string>(
+      /* id */ 'output',
+      /* name */ 'Output Type',
+      /* invokeSequence */ 'o',
+      /* longInvokeSequence */ 'output',
+      (node) => node.output.all.map((output) => output.type),
+      data,
+      filterSearchOptions
+    )
+
+    const nodeCategoryFilter = new NodeFilter<string>(
+      /* id */ 'category',
+      /* name */ 'Category',
+      /* invokeSequence */ 'c',
+      /* longInvokeSequence */ 'category',
+      (node) => [node.category],
+      data,
+      filterSearchOptions
+    )
+
+    const nodeSourceFilter = new NodeFilter<string>(
+      /* id */ 'source',
+      /* name */ 'Source',
+      /* invokeSequence */ 's',
+      /* longInvokeSequence */ 'source',
+      (node) => [node.nodeSource.displayText],
+      data,
+      filterSearchOptions
+    )
+
+    this.nodeFilters = [
+      inputTypeFilter,
+      outputTypeFilter,
+      nodeCategoryFilter,
+      nodeSourceFilter
+    ]
   }
 
   public endsWithFilterStartSequence(query: string): boolean {