From d2d2bc65074c515afabc743f5ceff060d7cb4ba2 Mon Sep 17 00:00:00 2001
From: Aaron <yangtao.yangtao@antgroup.com>
Date: Mon, 26 Feb 2024 16:56:43 +0800
Subject: [PATCH] refactor(elements): optimize elements types (#5472)

---
 packages/g6/src/elements/edges/base-edge.ts   | 36 ++++----
 packages/g6/src/elements/nodes/base-node.ts   | 88 ++++++++-----------
 packages/g6/src/elements/shapes/base-shape.ts | 14 +--
 packages/g6/src/runtime/element.ts            |  2 +-
 4 files changed, 61 insertions(+), 79 deletions(-)

diff --git a/packages/g6/src/elements/edges/base-edge.ts b/packages/g6/src/elements/edges/base-edge.ts
index 71ea64fd01d..62ba22f4819 100644
--- a/packages/g6/src/elements/edges/base-edge.ts
+++ b/packages/g6/src/elements/edges/base-edge.ts
@@ -29,23 +29,17 @@ import type { LabelStyleProps } from '../shapes';
 import { Label } from '../shapes';
 import { BaseShape } from '../shapes/base-shape';
 
-export type BaseEdgeStyleProps<StyleLifting extends keyof BaseEdgeKeyStyleProps = never> = BaseEdgeProps &
+export type BaseEdgeStyleProps = BaseEdgeProps &
   ShapeSwitch &
-  Pick<BaseEdgeKeyStyleProps, StyleLifting> &
   PrefixObject<EdgeLabelStyleProps, 'label'> &
   PrefixObject<PathStyleProps, 'halo'> &
   PrefixObject<EdgeArrowStyleProps, 'startArrow'> &
   PrefixObject<EdgeArrowStyleProps, 'endArrow'> &
   PrefixObject<LoopStyleProps, 'loop'>;
 
-type ParsedBaseEdgeStyleProps<StyleLifting extends keyof BaseEdgeKeyStyleProps = never> = Required<
-  BaseEdgeStyleProps<StyleLifting>
->;
-type BaseEdgeKeyStyleProps = PathStyleProps;
+type ParsedBaseEdgeStyleProps = Required<BaseEdgeStyleProps>;
 
-export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps = never> extends BaseShape<
-  BaseEdgeStyleProps<StyleLifting>
-> {
+export abstract class BaseEdge extends BaseShape<BaseEdgeStyleProps> {
   static defaultStyleProps: Partial<BaseEdgeStyleProps> = {
     isBillboard: true,
     label: true,
@@ -85,11 +79,11 @@ export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps
     loopClockwise: true,
   };
 
-  constructor(options: DisplayObjectConfig<BaseEdgeStyleProps<StyleLifting>>) {
+  constructor(options: DisplayObjectConfig<BaseEdgeStyleProps>) {
     super(deepMix({}, { style: BaseEdge.defaultStyleProps }, options));
   }
 
-  protected getKeyStyle(attributes: ParsedBaseEdgeStyleProps<StyleLifting>): PathStyleProps {
+  protected getKeyStyle(attributes: ParsedBaseEdgeStyleProps): PathStyleProps {
     const { sourceNode, targetNode, color, stroke, ...style } = this.getGraphicStyle(attributes);
 
     const path = isSameNode(sourceNode, targetNode) ? this.getLoopPath(attributes) : this.getKeyPath(attributes);
@@ -100,9 +94,9 @@ export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps
     };
   }
 
-  protected abstract getKeyPath(attributes: ParsedBaseEdgeStyleProps<StyleLifting>): PathArray;
+  protected abstract getKeyPath(attributes: ParsedBaseEdgeStyleProps): PathArray;
 
-  protected getLoopPath(attributes: ParsedBaseEdgeStyleProps<StyleLifting>): PathArray {
+  protected getLoopPath(attributes: ParsedBaseEdgeStyleProps): PathArray {
     const { sourceNode: node, sourcePort, targetPort } = attributes;
 
     const bbox = getNodeBBox(node);
@@ -117,7 +111,7 @@ export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps
     return getCubicLoopPath(node, position, clockwise, dist, sourcePort, targetPort);
   }
 
-  protected getEndpoints(attributes: ParsedBaseEdgeStyleProps<StyleLifting>): [Point, Point] {
+  protected getEndpoints(attributes: ParsedBaseEdgeStyleProps): [Point, Point] {
     const { sourceNode, targetNode, sourcePort: sourcePortKey, targetPort: targetPortKey } = attributes;
 
     const [sourcePort, targetPort] = findPorts(sourceNode, targetNode, sourcePortKey, targetPortKey);
@@ -128,7 +122,7 @@ export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps
     return [sourcePoint, targetPoint];
   }
 
-  protected getHaloStyle(attributes: ParsedBaseEdgeStyleProps<StyleLifting>): false | PathStyleProps {
+  protected getHaloStyle(attributes: ParsedBaseEdgeStyleProps): false | PathStyleProps {
     if (attributes.halo === false) return false;
 
     const keyStyle = this.getKeyStyle(attributes);
@@ -137,7 +131,7 @@ export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps
     return { ...keyStyle, ...haloStyle };
   }
 
-  protected getLabelStyle(attributes: ParsedBaseEdgeStyleProps<StyleLifting>): false | LabelStyleProps {
+  protected getLabelStyle(attributes: ParsedBaseEdgeStyleProps): false | LabelStyleProps {
     if (attributes.label === false || isEmpty(attributes.labelText)) return false;
 
     const labelStyle = subStyleProps<Required<EdgeLabelStyleProps>>(this.getGraphicStyle(attributes), 'label');
@@ -156,7 +150,7 @@ export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps
     return Object.assign({ wordWrapWidth }, labelPositionStyle, restStyle);
   }
 
-  protected drawArrow(attributes: ParsedBaseEdgeStyleProps<StyleLifting>, isStart: boolean) {
+  protected drawArrow(attributes: ParsedBaseEdgeStyleProps, isStart: boolean) {
     const arrowType = isStart ? 'startArrow' : 'endArrow';
     const arrowPresence = attributes[arrowType];
 
@@ -172,7 +166,7 @@ export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps
     }
   }
 
-  private getArrowStyle(attributes: ParsedBaseEdgeStyleProps<StyleLifting>, isStart: boolean) {
+  private getArrowStyle(attributes: ParsedBaseEdgeStyleProps, isStart: boolean) {
     const keyStyle = this.getKeyStyle(attributes) as BaseStyleProps;
     const arrowType = isStart ? 'startArrow' : 'endArrow';
     const { width, height, type, ctor, ...arrowStyle } = subStyleProps<Required<EdgeArrowStyleProps>>(
@@ -194,15 +188,15 @@ export abstract class BaseEdge<StyleLifting extends keyof BaseEdgeKeyStyleProps
     };
   }
 
-  protected drawLabelShape(attributes: ParsedBaseEdgeStyleProps<StyleLifting>, container: Group) {
+  protected drawLabelShape(attributes: ParsedBaseEdgeStyleProps, container: Group) {
     this.upsert('label', Label, this.getLabelStyle(attributes), container);
   }
 
-  protected drawHaloShape(attributes: ParsedBaseEdgeStyleProps<StyleLifting>, container: Group) {
+  protected drawHaloShape(attributes: ParsedBaseEdgeStyleProps, container: Group) {
     this.upsert('halo', Path, this.getHaloStyle(attributes), container);
   }
 
-  protected drawKeyShape(attributes: ParsedBaseEdgeStyleProps<StyleLifting>, container: Group): Path | undefined {
+  protected drawKeyShape(attributes: ParsedBaseEdgeStyleProps, container: Group): Path | undefined {
     return this.upsert('key', Path, this.getKeyStyle(attributes), container);
   }
 
diff --git a/packages/g6/src/elements/nodes/base-node.ts b/packages/g6/src/elements/nodes/base-node.ts
index bbfed6c80d5..c18aa767559 100644
--- a/packages/g6/src/elements/nodes/base-node.ts
+++ b/packages/g6/src/elements/nodes/base-node.ts
@@ -22,24 +22,32 @@ import { getWordWrapWidthByBox } from '../../utils/text';
 import type { BadgeStyleProps, IconStyleProps, LabelStyleProps } from '../shapes';
 import { Badge, BaseShape, Icon, Label } from '../shapes';
 
-export type BaseNodeStyleProps<
-  KeyStyleProps extends BaseStyleProps = BaseNodeProps,
-  StyleLifting extends keyof KeyStyleProps = never,
-> = BaseNodeProps &
-  ShapeSwitch &
-  // 将 KeyStyleProps 的属性提升到 BaseNodeStyleProps
-  // Lift the properties of KeyStyleProps to BaseNodeStyleProps
-  Pick<KeyStyleProps, StyleLifting> &
-  PrefixObject<NodeLabelStyleProps, 'label'> &
+export type BaseNodeStyleProps<KeyStyleProps extends BaseStyleProps = BaseNodeProps> = BaseNodeProps &
+  ShapeSwitch & {
+    /**
+     * <zh/> 连接桩
+     * <en/> Port
+     */
+    ports?: NodePortStyleProps[];
+    /**
+     * <zh/> 徽标
+     * <en/> Badge
+     */
+    badges?: NodeBadgeStyleProps[];
+    /**
+     * <zh/> 背景色板
+     * <en/> Background color palette
+     */
+    badgePalette?: string[] | CategoricalPalette;
+  } & PrefixObject<NodeLabelStyleProps, 'label'> &
   PrefixObject<KeyStyleProps, 'halo'> &
   PrefixObject<IconStyleProps, 'icon'> &
-  NodeBadgesStyleProps &
-  NodePortsStyleProps;
+  PrefixObject<BadgeStyleProps, 'badge'> &
+  PrefixObject<PortStyleProps, 'port'>;
 
-export type ParsedBaseNodeStyleProps<
-  KeyStyleProps extends BaseStyleProps,
-  StyleLifting extends keyof KeyStyleProps = never,
-> = Required<BaseNodeStyleProps<KeyStyleProps, StyleLifting>>;
+export type ParsedBaseNodeStyleProps<KeyStyleProps extends BaseStyleProps> = Required<
+  BaseNodeStyleProps<KeyStyleProps>
+>;
 
 /**
  * Design document: https://www.yuque.com/antv/g6/gl1iof1xpzg6ed98
@@ -47,11 +55,9 @@ export type ParsedBaseNodeStyleProps<
  * The P is the StyleProps of Key Shape.
  * The KeyShape is the type of the key shape.
  */
-export abstract class BaseNode<
-  KeyShape extends DisplayObject,
-  KeyStyleProps extends BaseStyleProps,
-  StyleLifting extends keyof KeyStyleProps = never,
-> extends BaseShape<BaseNodeStyleProps<KeyStyleProps, StyleLifting>> {
+export abstract class BaseNode<KeyShape extends DisplayObject, KeyStyleProps extends BaseStyleProps> extends BaseShape<
+  BaseNodeStyleProps<KeyStyleProps>
+> {
   static defaultStyleProps: BaseNodeStyleProps = {
     x: 0,
     y: 0,
@@ -79,7 +85,7 @@ export abstract class BaseNode<
     labelZIndex: 0,
   };
 
-  constructor(options: DisplayObjectConfig<BaseNodeStyleProps<KeyStyleProps, StyleLifting>>) {
+  constructor(options: DisplayObjectConfig<BaseNodeStyleProps<KeyStyleProps>>) {
     super(deepMix({}, { style: BaseNode.defaultStyleProps }, options));
   }
 
@@ -88,7 +94,7 @@ export abstract class BaseNode<
     return parseSize(size);
   }
 
-  protected getKeyStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>): KeyStyleProps {
+  protected getKeyStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>): KeyStyleProps {
     const { color, fill, ...style } = this.getGraphicStyle(attributes);
 
     return Object.assign(
@@ -97,7 +103,7 @@ export abstract class BaseNode<
     ) as unknown as KeyStyleProps;
   }
 
-  protected getLabelStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>): false | LabelStyleProps {
+  protected getLabelStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>): false | LabelStyleProps {
     if (attributes.label === false || isEmpty(attributes.labelText)) return false;
 
     const { position, maxWidth, ...labelStyle } = subStyleProps<Required<NodeLabelStyleProps>>(
@@ -114,7 +120,7 @@ export abstract class BaseNode<
     );
   }
 
-  protected getHaloStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>): false | KeyStyleProps {
+  protected getHaloStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>): false | KeyStyleProps {
     if (attributes.halo === false) return false;
 
     const { fill, ...keyStyle } = this.getKeyStyle(attributes);
@@ -123,7 +129,7 @@ export abstract class BaseNode<
     return { ...keyStyle, stroke: fill, ...haloStyle };
   }
 
-  protected getIconStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>): false | IconStyleProps {
+  protected getIconStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>): false | IconStyleProps {
     if (attributes.icon === false || isEmpty(attributes.iconText || attributes.iconSrc)) return false;
 
     const iconStyle = subStyleProps(this.getGraphicStyle(attributes), 'icon');
@@ -134,7 +140,7 @@ export abstract class BaseNode<
   }
 
   protected getBadgesStyle(
-    attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>,
+    attributes: ParsedBaseNodeStyleProps<KeyStyleProps>,
   ): Record<string, NodeBadgeStyleProps | false> {
     const badges = subObject(this.shapeMap, 'badge-');
     const badgesShapeStyle: Record<string, NodeBadgeStyleProps | false> = {};
@@ -166,9 +172,7 @@ export abstract class BaseNode<
     return { ...textStyle, ...restStyle };
   }
 
-  protected getPortsStyle(
-    attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>,
-  ): Record<string, PortStyleProps | false> {
+  protected getPortsStyle(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>): Record<string, PortStyleProps | false> {
     const ports = this.getPorts();
     const portsShapeStyle: Record<string, PortStyleProps | false> = {};
 
@@ -188,10 +192,7 @@ export abstract class BaseNode<
     return portsShapeStyle;
   }
 
-  protected getPortXY(
-    attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>,
-    style: NodePortStyleProps,
-  ): Point {
+  protected getPortXY(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>, style: NodePortStyleProps): Point {
     const { position = 'left' } = style;
     const bounds = this.getKey().getLocalBounds();
     return getPortPosition(bounds, position as PortPosition);
@@ -231,7 +232,7 @@ export abstract class BaseNode<
     return getRectIntersectPoint(point, keyShapeBounds);
   }
 
-  protected drawHaloShape(attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>, container: Group): void {
+  protected drawHaloShape(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>, container: Group): void {
     const keyShape = this.getKey();
     this.upsert(
       'halo',
@@ -241,14 +242,14 @@ export abstract class BaseNode<
     );
   }
 
-  protected drawBadgeShapes(attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>, container: Group): void {
+  protected drawBadgeShapes(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>, container: Group): void {
     const badgesStyle = this.getBadgesStyle(attributes);
     Object.keys(badgesStyle).forEach((key) => {
       this.upsert(`badge-${key}`, Badge, badgesStyle[key], container);
     });
   }
 
-  protected drawPortShapes(attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>, container: Group): void {
+  protected drawPortShapes(attributes: ParsedBaseNodeStyleProps<KeyStyleProps>, container: Group): void {
     const portsStyle = this.getPortsStyle(attributes);
     Object.keys(portsStyle).forEach((key) => {
       this.upsert(`port-${key}`, GCircle, portsStyle[key], container);
@@ -256,7 +257,7 @@ export abstract class BaseNode<
   }
 
   protected abstract drawKeyShape(
-    attributes: ParsedBaseNodeStyleProps<KeyStyleProps, StyleLifting>,
+    attributes: ParsedBaseNodeStyleProps<KeyStyleProps>,
     container: Group,
   ): KeyShape | undefined;
 
@@ -314,15 +315,6 @@ type NodeBadgeStyleProps = BadgeStyleProps & {
   position?: BadgePosition;
 };
 
-type NodeBadgesStyleProps = {
-  badges?: NodeBadgeStyleProps[];
-  /**
-   * <zh/> 背景色板
-   * <en/> Background color palette
-   */
-  badgePalette?: string[] | CategoricalPalette;
-} & PrefixObject<BadgeStyleProps, 'badge'>;
-
 export type NodePortStyleProps = Partial<PortStyleProps> & {
   /**
    * The key of the port. Default is the index of the port.
@@ -336,10 +328,6 @@ export type NodePortStyleProps = Partial<PortStyleProps> & {
   position: string | [number, number];
 };
 
-type NodePortsStyleProps = {
-  ports?: NodePortStyleProps[];
-} & PrefixObject<PortStyleProps, 'port'>;
-
 type ShapeSwitch = {
   label?: boolean;
   halo?: boolean;
diff --git a/packages/g6/src/elements/shapes/base-shape.ts b/packages/g6/src/elements/shapes/base-shape.ts
index 32da0db5a2e..f19e70d81f8 100644
--- a/packages/g6/src/elements/shapes/base-shape.ts
+++ b/packages/g6/src/elements/shapes/base-shape.ts
@@ -10,16 +10,16 @@ export interface BaseShapeStyleProps extends BaseStyleProps {
   y?: number | string;
 }
 
-export abstract class BaseShape<T extends BaseShapeStyleProps> extends CustomElement<T> {
-  constructor(options: DisplayObjectConfig<T>) {
+export abstract class BaseShape<StyleProps extends BaseShapeStyleProps> extends CustomElement<StyleProps> {
+  constructor(options: DisplayObjectConfig<StyleProps>) {
     super(options);
 
-    this.render(this.attributes as Required<T>, this);
+    this.render(this.attributes as Required<StyleProps>, this);
     this.bindEvents();
   }
 
   get parsedAttributes() {
-    return this.attributes as Required<T>;
+    return this.attributes as Required<StyleProps>;
   }
 
   /**
@@ -94,9 +94,9 @@ export abstract class BaseShape<T extends BaseShapeStyleProps> extends CustomEle
     return createAnimationsProxy(result, Object.values(this.animateMap));
   }
 
-  public update(attr: Partial<T> = {}): void {
+  public update(attr: Partial<StyleProps> = {}): void {
     this.attr(deepMix({}, this.attributes, attr));
-    return this.render(this.attributes as Required<T>, this);
+    return this.render(this.attributes as Required<StyleProps>, this);
   }
 
   /**
@@ -106,7 +106,7 @@ export abstract class BaseShape<T extends BaseShapeStyleProps> extends CustomEle
    * @param attributes
    * @param container
    */
-  public abstract render(attributes: Required<T>, container: Group): void;
+  public abstract render(attributes: Required<StyleProps>, container: Group): void;
 
   public bindEvents() {}
 
diff --git a/packages/g6/src/runtime/element.ts b/packages/g6/src/runtime/element.ts
index 002c4975d75..9fe4f8be670 100644
--- a/packages/g6/src/runtime/element.ts
+++ b/packages/g6/src/runtime/element.ts
@@ -321,7 +321,7 @@ export class ElementController {
   }
 
   public getEdges() {
-    return this.container.edge.children as BaseEdge<any>[];
+    return this.container.edge.children as BaseEdge[];
   }
 
   public getCombos() {