refactor: getPluginInstance allow match by type secondary (#6215)

Co-authored-by: antv <[email protected]>
antvis · Aug 22, 2024 · b0ce104 · b0ce104
1 parent b045322
commit b0ce104
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
diff --git a/packages/g6/__tests__/unit/runtime/graph/get-plugin-instantce.spec.ts b/packages/g6/__tests__/unit/runtime/graph/get-plugin-instantce.spec.ts
@@ -1,4 +1,4 @@
-import { BasePlugin, register } from '@/src';
+import { BasePlugin, ExtensionCategory, register } from '@/src';
 import { createGraph } from '@@/utils';
 
 describe('getPluginInstance', () => {
@@ -12,7 +12,7 @@ describe('getPluginInstance', () => {
       }
     }
 
-    register('plugin', 'custom', CustomPlugin);
+    register(ExtensionCategory.PLUGIN, 'custom', CustomPlugin);
     const graph = createGraph({
       plugins: [
         {
@@ -32,4 +32,26 @@ describe('getPluginInstance', () => {
     const undefinedPlugin = graph.getPluginInstance<CustomPlugin>('undefined-plugin');
     expect(undefinedPlugin).toBe(undefined);
   });
+
+  it('getPluginInstance by type', async () => {
+    const fn = jest.fn();
+
+    class CustomPlugin extends BasePlugin<any> {
+      api() {
+        fn();
+      }
+    }
+
+    register(ExtensionCategory.PLUGIN, 'custom-2', CustomPlugin);
+    const graph = createGraph({
+      plugins: ['custom-2', 'custom-2'],
+    });
+
+    await graph.draw();
+
+    const plugin = graph.getPluginInstance<CustomPlugin>('custom-2');
+    expect(plugin instanceof CustomPlugin).toBe(true);
+    plugin.api();
+    expect(fn).toHaveBeenCalled();
+  });
 });
diff --git a/packages/g6/src/runtime/plugin.ts b/packages/g6/src/runtime/plugin.ts
@@ -1,6 +1,7 @@
 import type { BasePlugin } from '../plugins/base-plugin';
 import { ExtensionController } from '../registry/extension';
 import type { CustomPluginOption, PluginOptions } from '../spec/plugin';
+import { print } from '../utils/print';
 import type { RuntimeContext } from './types';
 
 export class PluginController extends ExtensionController<BasePlugin<CustomPluginOption>> {
@@ -16,6 +17,12 @@ export class PluginController extends ExtensionController<BasePlugin<CustomPlugi
   }
 
   public getPluginInstance(key: string) {
-    return this.extensionMap[key];
+    const exactly = this.extensionMap[key];
+    if (exactly) return exactly;
+
+    print.warn(`Cannot find the plugin ${key}, will try to find it by type.`);
+
+    const fussily = this.extensions.find((extension) => extension.type === key);
+    if (fussily) return this.extensionMap[fussily.key];
   }
 }