feat: support shape type extension (#6209)

* feat: support shape extension * refactor: upsert support pass ctor type * refactor: fix types of extension 3d --------- Co-authored-by: antv <[email protected]>
antvis · Aug 22, 2024 · 1c9729b · 1c9729b
1 parent 9f36923
commit 1c9729b
Show file tree

Hide file tree

Showing 14 changed files with 102 additions and 31 deletions.
diff --git a/packages/g6-extension-3d/src/elements/base-node-3d.ts b/packages/g6-extension-3d/src/elements/base-node-3d.ts
@@ -51,9 +51,9 @@ export abstract class BaseNode3D<S extends BaseNode3DStyleProps> extends BaseNod
     return this.upsert('key', Mesh, this.getKeyStyle(attributes), container);
   }
 
-  protected abstract getGeometry(attributes: Required<S>): GGeometry<any> | undefined;
+  protected abstract getGeometry(attributes: Required<S>): GGeometry<any>;
 
-  protected getMaterial(attributes: Required<S>): GMaterial<any> | undefined {
+  protected getMaterial(attributes: Required<S>): GMaterial<any> {
     const { texture } = attributes;
     const materialStyle = subStyleProps<Material>(attributes, 'material');
     return createMaterial(this.plugin, materialStyle, texture);
@@ -63,6 +63,6 @@ export interface MeshStyleProps extends BaseStyleProps {
   x?: number | string;
   y?: number | string;
   z?: number | string;
-  geometry?: GGeometry<any>;
-  material?: GMaterial<any>;
+  geometry: GGeometry<any>;
+  material: GMaterial<any>;
 }
diff --git a/packages/g6-extension-3d/src/elements/capsule.ts b/packages/g6-extension-3d/src/elements/capsule.ts
@@ -20,7 +20,7 @@ export class Capsule extends BaseNode3D<CapsuleStyleProps> {
     super(deepMix({}, { style: Capsule.defaultStyleProps }, options));
   }
 
-  protected getGeometry(attributes: Required<CapsuleStyleProps>): GGeometry<any> | undefined {
+  protected getGeometry(attributes: Required<CapsuleStyleProps>): GGeometry<any> {
     const size = this.getSize();
     const { radius = size[0] / 2, height = size[1], heightSegments, sides } = attributes;
     return createGeometry('capsule', this.device, CapsuleGeometry, { radius, height, heightSegments, sides });

diff --git a/packages/g6-extension-3d/src/elements/cone.ts b/packages/g6-extension-3d/src/elements/cone.ts
@@ -20,7 +20,7 @@ export class Cone extends BaseNode3D<ConeStyleProps> {
     super(deepMix({}, { style: Cone.defaultStyleProps }, options));
   }
 
-  protected getGeometry(attributes: Required<ConeStyleProps>): GGeometry<any> | undefined {
+  protected getGeometry(attributes: Required<ConeStyleProps>): GGeometry<any> {
     const size = this.getSize();
     const {
       baseRadius = size[0] / 2,

diff --git a/packages/g6-extension-3d/src/elements/cube.ts b/packages/g6-extension-3d/src/elements/cube.ts
@@ -19,7 +19,7 @@ export class Cube extends BaseNode3D<CubeStyleProps> {
     super(deepMix({}, { style: Cube.defaultStyleProps }, options));
   }
 
-  protected getGeometry(attributes: Required<CubeStyleProps>): GGeometry<any> | undefined {
+  protected getGeometry(attributes: Required<CubeStyleProps>): GGeometry<any> {
     const size = this.getSize();
     const {
       width = size[0],

diff --git a/packages/g6-extension-3d/src/elements/cylinder.ts b/packages/g6-extension-3d/src/elements/cylinder.ts
@@ -20,7 +20,7 @@ export class Cylinder extends BaseNode3D<CylinderStyleProps> {
     super(deepMix({}, { style: Cylinder.defaultStyleProps }, options));
   }
 
-  protected getGeometry(attributes: Required<CylinderStyleProps>): GGeometry<any> | undefined {
+  protected getGeometry(attributes: Required<CylinderStyleProps>): GGeometry<any> {
     const size = this.getSize();
     const { radius = size[0] / 2, height = size[1], heightSegments, capSegments } = attributes;
     return createGeometry('cylinder', this.device, CylinderGeometry, { radius, height, heightSegments, capSegments });

diff --git a/packages/g6-extension-3d/src/elements/plane.ts b/packages/g6-extension-3d/src/elements/plane.ts
@@ -17,7 +17,7 @@ export class Plane extends BaseNode3D<PlaneStyleProps> {
     super(deepMix({}, { style: Plane.defaultStyleProps }, options));
   }
 
-  protected getGeometry(attributes: Required<PlaneStyleProps>): GGeometry<any> | undefined {
+  protected getGeometry(attributes: Required<PlaneStyleProps>): GGeometry<any> {
     const size = this.getSize();
     const { width = size[0], depth = size[1], widthSegments, depthSegments } = attributes;
     return createGeometry('plane', this.device, PlaneGeometry, { width, depth, widthSegments, depthSegments });

diff --git a/packages/g6-extension-3d/src/elements/sphere.ts b/packages/g6-extension-3d/src/elements/sphere.ts
@@ -20,7 +20,7 @@ export class Sphere extends BaseNode3D<SphereStyleProps> {
     super(deepMix({}, { style: Sphere.defaultStyleProps }, options));
   }
 
-  protected getGeometry(attributes: Required<SphereStyleProps>): GGeometry<any> | undefined {
+  protected getGeometry(attributes: Required<SphereStyleProps>): GGeometry<any> {
     const size = this.getSize();
     const { radius = size[0] / 2, latitudeBands, longitudeBands } = attributes;
     return createGeometry('sphere', this.device, SphereGeometry, { radius, latitudeBands, longitudeBands });

diff --git a/packages/g6-extension-3d/src/elements/torus.ts b/packages/g6-extension-3d/src/elements/torus.ts
@@ -20,7 +20,7 @@ export class Torus extends BaseNode3D<TorusStyleProps> {
     super(deepMix({}, { style: Torus.defaultStyleProps }, options));
   }
 
-  protected getGeometry(attributes: Required<TorusStyleProps>): GGeometry<any> | undefined {
+  protected getGeometry(attributes: Required<TorusStyleProps>): GGeometry<any> {
     const size = this.getSize();
     const { tubeRadius = size[0] / 2, ringRadius = size[1] / 2, segments, sides } = attributes;
     return createGeometry('torus', this.device, TorusGeometry, { tubeRadius, ringRadius, segments, sides });

diff --git a/packages/g6/__tests__/unit/registry.spec.ts b/packages/g6/__tests__/unit/registry.spec.ts
@@ -7,6 +7,7 @@ import {
   Diamond,
   Donut,
   Ellipse,
+  ExtensionCategory,
   HTML,
   Hexagon,
   Image,
@@ -22,11 +23,12 @@ import {
   register,
 } from '@/src';
 import { dark, light } from '@/src/themes';
+import { Circle as GCircle } from '@antv/g';
 import { pick } from '@antv/util';
 
 describe('registry', () => {
   it('registerBuiltInPlugins', () => {
-    expect(getExtensions('node')).toEqual({
+    expect(getExtensions(ExtensionCategory.NODE)).toEqual({
       circle: Circle,
       ellipse: Ellipse,
       image: Image,
@@ -38,19 +40,19 @@ describe('registry', () => {
       hexagon: Hexagon,
       html: HTML,
     });
-    expect(getExtensions('edge')).toEqual({
+    expect(getExtensions(ExtensionCategory.EDGE)).toEqual({
       cubic: Cubic,
       line: Line,
       polyline: Polyline,
       quadratic: Quadratic,
       'cubic-horizontal': CubicHorizontal,
       'cubic-vertical': CubicVertical,
     });
-    expect(getExtensions('combo')).toEqual({
+    expect(getExtensions(ExtensionCategory.COMBO)).toEqual({
       circle: CircleCombo,
       rect: RectCombo,
     });
-    expect(getExtensions('theme')).toEqual({
+    expect(getExtensions(ExtensionCategory.THEME)).toEqual({
       dark,
       light,
     });
@@ -60,22 +62,22 @@ describe('registry', () => {
     class CircleNode {}
     class RectNode {}
     class Edge {}
-    register('node', 'circle-node', CircleNode as any);
-    register('node', 'rect-node', RectNode as any);
-    register('edge', 'line-edge', Edge as any);
-    expect(getExtension('node', 'circle-node')).toEqual(CircleNode);
-    expect(getExtension('node', 'rect-node')).toEqual(RectNode);
-    expect(getExtension('node', 'diamond-node')).toEqual(undefined);
-    expect(getExtension('edge', 'line-edge')).toEqual(Edge);
+    register(ExtensionCategory.NODE, 'circle-node', CircleNode as any);
+    register(ExtensionCategory.NODE, 'rect-node', RectNode as any);
+    register(ExtensionCategory.EDGE, 'line-edge', Edge as any);
+    expect(getExtension(ExtensionCategory.NODE, 'circle-node')).toEqual(CircleNode);
+    expect(getExtension(ExtensionCategory.NODE, 'rect-node')).toEqual(RectNode);
+    expect(getExtension(ExtensionCategory.NODE, 'diamond-node')).toEqual(undefined);
+    expect(getExtension(ExtensionCategory.EDGE, 'line-edge')).toEqual(Edge);
 
     const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation();
 
-    register('node', 'circle-node', CircleNode as any);
+    register(ExtensionCategory.NODE, 'circle-node', CircleNode as any);
     expect(consoleErrorSpy).toHaveBeenCalledTimes(0);
 
     consoleErrorSpy.mockRestore();
 
-    expect(pick(getExtensions('node'), ['circle-node', 'rect-node'])).toEqual({
+    expect(pick(getExtensions(ExtensionCategory.NODE), ['circle-node', 'rect-node'])).toEqual({
       'circle-node': CircleNode,
       'rect-node': RectNode,
     });
@@ -84,8 +86,31 @@ describe('registry', () => {
   it('override', () => {
     class CircleNode {}
     class RectNode {}
-    register('node', 'circle-node', CircleNode as any);
-    register('node', 'circle-node', RectNode as any, true);
-    expect(getExtension('node', 'circle-node')).toEqual(RectNode);
+    register(ExtensionCategory.NODE, 'circle-node', CircleNode as any);
+    register(ExtensionCategory.NODE, 'circle-node', RectNode as any, true);
+    expect(getExtension(ExtensionCategory.NODE, 'circle-node')).toEqual(RectNode);
+  });
+
+  it('register shape', () => {
+    const shapes = getExtensions(ExtensionCategory.SHAPE);
+    expect(Object.keys(shapes)).toEqual([
+      'circle',
+      'ellipse',
+      'group',
+      'html',
+      'image',
+      'line',
+      'path',
+      'polygon',
+      'polyline',
+      'rect',
+      'text',
+      'label',
+      'badge',
+    ]);
+
+    register(ExtensionCategory.SHAPE, 'circle-shape', GCircle);
+
+    expect(getExtension(ExtensionCategory.SHAPE, 'circle-shape')).toEqual(GCircle);
   });
 });
diff --git a/packages/g6/src/constants/registry.ts b/packages/g6/src/constants/registry.ts
@@ -59,4 +59,10 @@ export enum ExtensionCategory {
    * <en/> Data transform
    */
   TRANSFORM = 'transform',
+  /**
+   * <zh/> 图形
+   *
+   * <en/> Shape
+   */
+  SHAPE = 'shape',
 }
diff --git a/packages/g6/src/elements/shapes/base-shape.ts b/packages/g6/src/elements/shapes/base-shape.ts
@@ -1,13 +1,16 @@
 import type { BaseStyleProps, DisplayObject, DisplayObjectConfig, Group, IAnimation } from '@antv/g';
 import { CustomElement } from '@antv/g';
 import { isEmpty, isFunction, upperFirst } from '@antv/util';
+import { ExtensionCategory } from '../../constants';
 import type { Keyframe } from '../../types';
 import { createAnimationsProxy, preprocessKeyframes } from '../../utils/animation';
 import { updateStyle } from '../../utils/element';
 import { subObject } from '../../utils/prefix';
+import { format } from '../../utils/print';
 import { getSubShapeStyle } from '../../utils/style';
 import { replaceTranslateInTransform } from '../../utils/transform';
 import { setVisibility } from '../../utils/visibility';
+import { getExtension } from './../../registry/get';
 
 export interface BaseShapeStyleProps extends BaseStyleProps {}
 
@@ -51,7 +54,7 @@ export abstract class BaseShape<StyleProps extends BaseShapeStyleProps> extends
    */
   protected upsert<T extends DisplayObject>(
     className: string,
-    Ctor: { new (...args: any[]): T },
+    Ctor: string | { new (...args: any[]): T },
     style: T['attributes'] | false,
     container: DisplayObject,
     hooks?: UpsertHooks,
@@ -69,20 +72,26 @@ export abstract class BaseShape<StyleProps extends BaseShapeStyleProps> extends
       return;
     }
 
+    const _Ctor = typeof Ctor === 'string' ? getExtension(ExtensionCategory.SHAPE, Ctor) : Ctor;
+
+    if (!_Ctor) {
+      throw new Error(format(`Shape ${Ctor} not found`));
+    }
+
     // create
-    if (!target || target.destroyed || !(target instanceof Ctor)) {
+    if (!target || target.destroyed || !(target instanceof _Ctor)) {
       if (target) {
         hooks?.beforeDestroy?.(target);
         target?.destroy();
         hooks?.afterDestroy?.(target);
       }
 
       hooks?.beforeCreate?.();
-      const instance = new Ctor({ className, style });
+      const instance = new _Ctor({ className, style });
       container.appendChild(instance);
       this.shapeMap[className] = instance;
       hooks?.afterCreate?.(instance);
-      return instance;
+      return instance as T;
     }
 
     // update

diff --git a/packages/g6/src/registry/build-in.ts b/packages/g6/src/registry/build-in.ts
@@ -1,3 +1,15 @@
+import {
+  Circle as GCircle,
+  Ellipse as GEllipse,
+  Group as GGroup,
+  HTML as GHTML,
+  Line as GLine,
+  Path as GPath,
+  Polygon as GPolygon,
+  Polyline as GPolyline,
+  Rect as GRect,
+  Text as GText,
+} from '@antv/g';
 import { ComboCollapse, ComboExpand, Fade, NodeCollapse, NodeExpand, PathIn, PathOut, Translate } from '../animations';
 import {
   BrushSelect,
@@ -35,6 +47,7 @@ import {
   Star,
   Triangle,
 } from '../elements';
+import { Badge as BadgeShape, Image as ImageShape, Label as LabelShape } from '../elements/shapes';
 import {
   AntVDagreLayout,
   CircularLayout,
@@ -193,6 +206,21 @@ const BUILT_IN_EXTENSIONS: ExtensionRegistry = {
     'process-parallel-edges': ProcessParallelEdges,
     'get-edge-actual-ends': GetEdgeActualEnds,
   },
+  shape: {
+    circle: GCircle,
+    ellipse: GEllipse,
+    group: GGroup,
+    html: GHTML,
+    image: ImageShape,
+    line: GLine,
+    path: GPath,
+    polygon: GPolygon,
+    polyline: GPolyline,
+    rect: GRect,
+    text: GText,
+    label: LabelShape,
+    badge: BadgeShape,
+  },
 };
 
 import type { ExtensionCategory } from '../constants';

diff --git a/packages/g6/src/registry/store.ts b/packages/g6/src/registry/store.ts
@@ -16,4 +16,5 @@ export const EXTENSION_REGISTRY: ExtensionRegistry = {
   theme: {},
   plugin: {},
   transform: {},
+  shape: {},
 };
diff --git a/packages/g6/src/registry/types.ts b/packages/g6/src/registry/types.ts
@@ -1,3 +1,4 @@
+import type { DisplayObject } from '@antv/g';
 import type { STDAnimation } from '../animations/types';
 import type { Behavior } from '../behaviors/types';
 import type { Layout } from '../layouts/types';
@@ -23,4 +24,5 @@ export interface ExtensionRegistry {
   plugin: Record<string, { new (...args: any[]): Plugin }>;
   animation: Record<string, STDAnimation>; // animation spec
   transform: Record<string, { new (...args: any[]): Transform }>;
+  shape: Record<string, { new (...args: any[]): DisplayObject }>;
 }