feat: adds support for generic and inferred typing based on graph data

packages/g6/src/elements/combos/base-combo.ts

@@ -23,14 +23,15 @@ import type { BaseNodeStyleProps } from '../nodes';
 import { BaseNode } from '../nodes';
 import { Icon, IconStyleProps } from '../shapes';
 import { connectImage, dispatchPositionChange } from '../shapes/image';
+import { UnknownStruct } from '../../types/utility';
 
 /**
  * <zh/> 组合通用样式配置项
  *
  * <en/> Common style props for combo
  */
-export interface BaseComboStyleProps
-  extends BaseNodeStyleProps,
+export interface BaseComboStyleProps<NodeType extends UnknownStruct = UnknownStruct, ComboType extends UnknownStruct = UnknownStruct>
+  extends Omit<BaseNodeStyleProps<NodeType>, 'childrenData'>,
     Prefix<'collapsed', BaseStyleProps>,
     Prefix<'collapsedMarker', CollapsedMarkerStyleProps> {
   /**
@@ -60,7 +61,7 @@ export interface BaseComboStyleProps
    *
    * <en/> If the combo is collapsed, children may be empty, and the complete child element data can be obtained through childrenData
    */
-  childrenData?: NodeLikeData[];
+  childrenData?: NodeLikeData<NodeType, ComboType>[];
   /**
    * <zh/> 组合的内边距，只在展开状态下生效
    *

packages/g6/src/elements/combos/circle.ts

@@ -7,29 +7,30 @@ import { subStyleProps } from '../../utils/prefix';
 import { parseSize } from '../../utils/size';
 import type { BaseComboStyleProps } from './base-combo';
 import { BaseCombo } from './base-combo';
+import { UnknownStruct } from '../../types/utility';
 
 /**
  * <zh/> 圆形组合样式配置项
  *
  * <en/> Circle combo style props
  */
-export interface CircleComboStyleProps extends BaseComboStyleProps {}
+export interface CircleComboStyleProps<NodeType extends UnknownStruct = UnknownStruct, ComboType extends UnknownStruct = UnknownStruct> extends BaseComboStyleProps<NodeType, ComboType> {}
 
 /**
  * <zh/> 圆形组合
  *
  * <en/> Circle combo
  */
-export class CircleCombo extends BaseCombo<CircleComboStyleProps> {
-  constructor(options: DisplayObjectConfig<CircleComboStyleProps>) {
+export class CircleCombo<NodeType extends UnknownStruct = UnknownStruct, ComboType extends UnknownStruct = UnknownStruct> extends BaseCombo<CircleComboStyleProps<NodeType, ComboType>> {
+  constructor(options: DisplayObjectConfig<CircleComboStyleProps<NodeType, ComboType>>) {
     super(options);
   }
 
-  protected drawKeyShape(attributes: Required<CircleComboStyleProps>, container: Group): GCircle | undefined {
+  protected drawKeyShape(attributes: Required<CircleComboStyleProps<NodeType, ComboType>>, container: Group): GCircle | undefined {
     return this.upsert('key', GCircle, this.getKeyStyle(attributes), container);
   }
 
-  protected getKeyStyle(attributes: Required<CircleComboStyleProps>): GCircleStyleProps {
+  protected getKeyStyle(attributes: Required<CircleComboStyleProps<NodeType, ComboType>>): GCircleStyleProps {
     const { collapsed } = attributes;
     const keyStyle = super.getKeyStyle(attributes);
 
@@ -41,13 +42,13 @@ export class CircleCombo extends BaseCombo<CircleComboStyleProps> {
     };
   }
 
-  protected getCollapsedKeySize(attributes: Required<CircleComboStyleProps>): STDSize {
+  protected getCollapsedKeySize(attributes: Required<CircleComboStyleProps<NodeType, ComboType>>): STDSize {
     const [collapsedWidth, collapsedHeight] = parseSize(attributes.collapsedSize);
     const collapsedR = Math.max(collapsedWidth, collapsedHeight) / 2;
     return [collapsedR * 2, collapsedR * 2, 0];
   }
 
-  protected getExpandedKeySize(attributes: Required<CircleComboStyleProps>): STDSize {
+  protected getExpandedKeySize(attributes: Required<CircleComboStyleProps<NodeType, ComboType>>): STDSize {
     const contentBBox = this.getContentBBox(attributes);
     const [width, height] = getBBoxSize(contentBBox);
     const expandedR = Math.sqrt(width ** 2 + height ** 2) / 2;

packages/g6/src/elements/combos/rect.ts

@@ -3,29 +3,30 @@ import { Rect as GRect, Group } from '@antv/g';
 import { subStyleProps } from '../../utils/prefix';
 import type { BaseComboStyleProps } from './base-combo';
 import { BaseCombo } from './base-combo';
+import { UnknownStruct } from '../../types/utility';
 
 /**
  * <zh/> 矩形组合样式配置项
  *
  * <en/> Rect combo style props
  */
-export interface RectComboStyleProps extends BaseComboStyleProps {}
+export interface RectComboStyleProps<NodeType extends UnknownStruct = UnknownStruct, ComboType extends UnknownStruct = UnknownStruct> extends BaseComboStyleProps<NodeType, ComboType> {}
 
 /**
  * <zh/> 矩形组合
  *
  * <en/> Rect combo
  */
-export class RectCombo extends BaseCombo<RectComboStyleProps> {
-  constructor(options: DisplayObjectConfig<RectComboStyleProps>) {
+export class RectCombo<NodeType extends UnknownStruct = UnknownStruct, ComboType extends UnknownStruct = UnknownStruct> extends BaseCombo<RectComboStyleProps<NodeType, ComboType>> {
+  constructor(options: DisplayObjectConfig<RectComboStyleProps<NodeType, ComboType>>) {
     super(options);
   }
 
-  protected drawKeyShape(attributes: Required<RectComboStyleProps>, container: Group): GRect | undefined {
+  protected drawKeyShape(attributes: Required<RectComboStyleProps<NodeType, ComboType>>, container: Group): GRect | undefined {
     return this.upsert('key', GRect, this.getKeyStyle(attributes), container);
   }
 
-  protected getKeyStyle(attributes: Required<RectComboStyleProps>): GRectStyleProps {
+  protected getKeyStyle(attributes: Required<RectComboStyleProps<NodeType, ComboType>>): GRectStyleProps {
     const keyStyle = super.getKeyStyle(attributes);
 
     const [width, height] = this.getKeySize(attributes);

packages/g6/src/elements/nodes/base-node.ts

@@ -29,13 +29,15 @@ import { BaseElement } from '../base-element';
 import type { BadgeStyleProps, BaseShapeStyleProps, IconStyleProps, LabelStyleProps } from '../shapes';
 import { Badge, Icon, Label } from '../shapes';
 import { connectImage, dispatchPositionChange } from '../shapes/image';
+import { UnknownStruct } from '../../types/utility';
 
 /**
  * <zh/> 节点通用样式配置项
  *
  * <en/> Base node style props
  */
-export interface BaseNodeStyleProps
+export interface BaseNodeStyleProps<
+  NodeType extends UnknownStruct = UnknownStruct>
   extends BaseShapeStyleProps,
     Prefix<'label', NodeLabelStyleProps>,
     Prefix<'halo', BaseStyleProps>,
@@ -97,7 +99,7 @@ export interface BaseNodeStyleProps
    * <en/> Only valid in the tree graph. If the current node is collapsed, children may be empty, and the complete child element data can be obtained through childrenData
    * @ignore
    */
-  childrenData?: NodeData[];
+  childrenData?: NodeData<NodeType>[];
   /**
    * <zh/> 是否显示节点标签
    *

packages/g6/src/runtime/data.ts

@@ -29,9 +29,10 @@ import { positionOf } from '../utils/position';
 import { format, print } from '../utils/print';
 import { dfs } from '../utils/traverse';
 import { add } from '../utils/vector';
+import { InferGraphDataTypes } from '../spec/graph';
 
-export class DataController {
-  public model: GraphLib<NodeLikeData, EdgeData>;
+export class DataController<D extends GraphData = GraphData> {
+  public model: GraphLib<NodeLikeData<InferGraphDataTypes<D>['node'], InferGraphDataTypes<D>['combo']>, EdgeData<InferGraphDataTypes<D>['edge']>>;
 
   /**
    * <zh/> 最近一次删除的 combo 的 id
@@ -132,37 +133,37 @@ export class DataController {
   }
 
   public getNodeData(ids?: ID[]) {
-    return this.model.getAllNodes().reduce((acc, node) => {
-      const data = toG6Data(node);
+    return this.model.getAllNodes().reduce<NodeData<InferGraphDataTypes<D>['node']>[]>((acc, node) => {
+      const data = toG6Data(node) as NodeData<InferGraphDataTypes<D>['node']>;
       if (this.isCombo(idOf(data))) return acc;
       if (ids === undefined) acc.push(data);
       else ids.includes(idOf(data)) && acc.push(data);
       return acc;
-    }, [] as NodeData[]);
+    }, []);
   }
 
   public getEdgeDatum(id: ID) {
     return toG6Data(this.model.getEdge(id));
   }
 
   public getEdgeData(ids?: ID[]) {
-    return this.model.getAllEdges().reduce((acc, edge) => {
+    return this.model.getAllEdges().reduce<EdgeData<InferGraphDataTypes<D>['edge']>[]>((acc, edge) => {
       const data = toG6Data(edge);
       if (ids === undefined) acc.push(data);
       else ids.includes(idOf(data)) && acc.push(data);
       return acc;
-    }, [] as EdgeData[]);
+    }, []);
   }
 
   public getComboData(ids?: ID[]) {
-    return this.model.getAllNodes().reduce((acc, combo) => {
+    return this.model.getAllNodes().reduce<ComboData<InferGraphDataTypes<D>['combo']>[]>((acc, combo) => {
       const data = toG6Data(combo);
       if (!this.isCombo(idOf(data))) return acc;
 
-      if (ids === undefined) acc.push(data as ComboData);
-      else ids.includes(idOf(data)) && acc.push(data as ComboData);
+      if (ids === undefined) acc.push(data);
+      else ids.includes(idOf(data)) && acc.push(data);
       return acc;
-    }, [] as ComboData[]);
+    }, []);
   }
 
   public getRootsData(hierarchyKey: HierarchyKey = TREE_KEY) {
@@ -200,11 +201,11 @@ export class DataController {
     return parent ? toG6Data(parent) : undefined;
   }
 
-  public getChildrenData(id: ID): NodeLikeData[] {
+  public getChildrenData(id: ID): NodeLikeData<InferGraphDataTypes<D>['node'], InferGraphDataTypes<D>['combo']>[] {
     const structureKey = this.getElementType(id) === 'node' ? TREE_KEY : COMBO_KEY;
     const { model } = this;
     if (!model.hasNode(id) || !model.hasTreeStructure(structureKey)) return [];
-    return model.getChildren(id, structureKey).map(toG6Data);
+    return model.getChildren(id, structureKey).map(toG6Data) as NodeLikeData<InferGraphDataTypes<D>['node'], InferGraphDataTypes<D>['combo']>[];
   }
 
   /**
@@ -336,7 +337,7 @@ export class DataController {
     this.computeZIndex(data, 'add');
   }
 
-  public addNodeData(nodes: NodeData[] = []) {
+  public addNodeData(nodes: NodeData<InferGraphDataTypes<D>['node']>[] = []) {
     if (!nodes.length) return;
     this.model.addNodes(
       nodes.map((node) => {
@@ -349,7 +350,7 @@ export class DataController {
     this.computeZIndex({ nodes }, 'add');
   }
 
-  public addEdgeData(edges: EdgeData[] = []) {
+  public addEdgeData(edges: EdgeData<InferGraphDataTypes<D>['edge']>[] = []) {
     if (!edges.length) return;
     this.model.addEdges(
       edges.map((edge) => {
@@ -361,7 +362,7 @@ export class DataController {
     this.computeZIndex({ edges }, 'add');
   }
 
-  public addComboData(combos: ComboData[] = []) {
+  public addComboData(combos: ComboData<InferGraphDataTypes<D>['combo']>[] = []) {
     if (!combos.length) return;
     const { model } = this;
 
@@ -382,8 +383,8 @@ export class DataController {
     this.computeZIndex({ combos }, 'add');
   }
 
-  public addChildrenData(parentId: ID, childrenData: NodeData[]) {
-    const parentData = this.getNodeLikeDatum(parentId) as NodeData;
+  public addChildrenData(parentId: ID, childrenData: NodeData<InferGraphDataTypes<D>['node']>[]) {
+    const parentData = this.getNodeLikeDatum(parentId) as NodeData<InferGraphDataTypes<D>['node']>;
     const childrenId = childrenData.map(idOf);
     this.addNodeData(childrenData);
     this.updateNodeData([{ id: parentId, children: [...(parentData.children || []), ...childrenId] }]);
@@ -562,17 +563,17 @@ export class DataController {
     this.computeZIndex(data, 'update');
   }
 
-  public updateNodeData(nodes: PartialNodeLikeData<NodeData>[] = []) {
+  public updateNodeData(nodes: PartialNodeLikeData<NodeData<InferGraphDataTypes<D>['node']>>[] = []) {
     if (!nodes.length) return;
     const { model } = this;
     this.batch(() => {
-      const modifiedNodes: NodeData[] = [];
+      const modifiedNodes: NodeData<InferGraphDataTypes<D>['node']>[] = [];
       nodes.forEach((modifiedNode) => {
         const id = idOf(modifiedNode);
-        const originalNode = toG6Data(model.getNode(id));
+        const originalNode = toG6Data(model.getNode(id)) as NodeData<InferGraphDataTypes<D>['node']>;
         if (isElementDataEqual(originalNode, modifiedNode)) return;
 
-        const value = mergeElementsData(originalNode, modifiedNode);
+        const value = mergeElementsData<D, NodeData<InferGraphDataTypes<D>['node']>>(originalNode, modifiedNode);
         this.pushChange({ value, original: originalNode, type: ChangeType.NodeUpdated });
         model.mergeNodeData(id, value);
         modifiedNodes.push(value);
@@ -787,12 +788,12 @@ export class DataController {
     const dy = ty - comboY;
     const dz = tz - comboZ;
 
-    dfs<NodeLikeData>(
+    dfs<NodeLikeData<InferGraphDataTypes<D>['node'], InferGraphDataTypes<D>['combo']>>(
       combo,
       (succeed) => {
         const succeedId = idOf(succeed);
         const [x, y, z] = positionOf(succeed);
-        const value = mergeElementsData(succeed, {
+        const value = mergeElementsData<D, NodeData<InferGraphDataTypes<D>['node']>>(succeed, {
           style: { x: x + dx, y: y + dy, z: z + dz },
         });
         this.pushChange({
@@ -875,7 +876,7 @@ export class DataController {
         const childData = toG6Data(child);
         const childId = idOf(childData);
         this.setParent(idOf(childData), grandParent, COMBO_KEY, false);
-        const value = mergeElementsData(childData, {
+        const value = mergeElementsData<D, NodeData<InferGraphDataTypes<D>['node']>>(childData, {
           id: idOf(childData),
           combo: grandParent,
         });

packages/g6/src/runtime/graph.ts

@@ -62,9 +62,10 @@ import { PluginController } from './plugin';
 import { TransformController } from './transform';
 import { RuntimeContext } from './types';
 import { ViewportController } from './viewport';
+import { InferGraphDataTypes } from '../spec/graph';
 
-export class Graph extends EventEmitter {
-  private options: GraphOptions = {};
+export class Graph<D extends GraphData = GraphData> extends EventEmitter {
+  private options: GraphOptions<D> = {};
 
   /**
    * @internal
@@ -96,7 +97,7 @@ export class Graph extends EventEmitter {
     model: new DataController(),
   };
 
-  constructor(options: GraphOptions) {
+  constructor(options: GraphOptions<D>) {
     super();
     this._setOptions(Object.assign({}, Graph.defaultOptions, options), true);
     this.context.graph = this;
@@ -112,7 +113,7 @@ export class Graph extends EventEmitter {
    * @returns <zh/> 配置项 | <en/> options
    * @apiCategory option
    */
-  public getOptions(): GraphOptions {
+  public getOptions(): GraphOptions<D> {
     return this.options;
   }
 
@@ -127,11 +128,11 @@ export class Graph extends EventEmitter {
    * <en/> To update devicePixelRatio and container properties, please destroy and recreate the instance
    * @apiCategory option
    */
-  public setOptions(options: GraphOptions): void {
+  public setOptions(options: GraphOptions<D>): void {
     this._setOptions(options, false);
   }
 
-  private _setOptions(options: GraphOptions, isInit: boolean) {
+  private _setOptions(options: GraphOptions<D>, isInit: boolean) {
     this.updateCanvas(options);
     Object.assign(this.options, inferOptions(options));
 
@@ -186,7 +187,7 @@ export class Graph extends EventEmitter {
    * @param zoomRange - <zh/> 缩放区间 | <en/> zoom range
    * @apiCategory viewport
    */
-  public setZoomRange(zoomRange: GraphOptions['zoomRange']): void {
+  public setZoomRange(zoomRange: GraphOptions<D>['zoomRange']): void {
     this.options.zoomRange = zoomRange;
   }
 
@@ -197,7 +198,7 @@ export class Graph extends EventEmitter {
    * @returns <zh/> 缩放区间 | <en/> zoom range
    * @apiCategory viewport
    */
-  public getZoomRange(): GraphOptions['zoomRange'] {
+  public getZoomRange(): GraphOptions<D>['zoomRange'] {
     return this.options.zoomRange;
   }
 
@@ -244,7 +245,7 @@ export class Graph extends EventEmitter {
    * <en/> The value of `options.combo`
    * @apiCategory element
    */
-  public setCombo(combo: ComboOptions): void {
+  public setCombo(combo: ComboOptions<InferGraphDataTypes<D>['node'], InferGraphDataTypes<D>['combo']>): void {
     this.options.combo = combo;
     this.context.model.refreshData();
   }

packages/g6/src/runtime/options.ts

@@ -1,4 +1,4 @@
-import type { GraphOptions } from '../spec';
+import type { GraphData, GraphOptions } from '../spec';
 
 /**
  * <zh/> 基于用户传入的配置，推断出最终的配置
@@ -7,7 +7,7 @@ import type { GraphOptions } from '../spec';
  * @param options - <zh/> 用户传入的配置 | <en/> Configuration passed by the user
  * @returns <zh/> 最终的配置 | <en/> Final configuration
  */
-export function inferOptions(options: GraphOptions): GraphOptions {
+export function inferOptions<D extends GraphData = GraphData>(options: GraphOptions<D>): GraphOptions<D> {
   const flow = [inferLayoutOptions];
   return flow.reduce((acc, infer) => infer(acc), options);
 }
@@ -19,7 +19,7 @@ export function inferOptions(options: GraphOptions): GraphOptions {
  * @param options - <zh/> 用户传入的配置 | <en/> Configuration passed by the user
  * @returns <zh/> 最终的配置 | <en/> Final configuration
  */
-function inferLayoutOptions(options: GraphOptions): GraphOptions {
+function inferLayoutOptions<D extends GraphData = GraphData>(options: GraphOptions<D>): GraphOptions<D> {
   if (!options.layout) return options;
   if (Array.isArray(options.layout)) return options;
   if ('preLayout' in options.layout) return options;