refactor: collapse-expand-react-node transform

packages/graphs/src/core/transform/collapse-expand-react-node.tsx

@@ -1,4 +1,4 @@
-import type { BaseTransformOptions, CardinalPlacement, Graph, NodeData, RuntimeContext } from '@antv/g6';
+import type { BaseTransformOptions, CardinalPlacement, DrawData, Graph, NodeData, RuntimeContext } from '@antv/g6';
 import { BaseTransform, idOf } from '@antv/g6';
 import { get, has, set } from 'lodash';
 import React from 'react';
@@ -84,21 +84,21 @@ export class CollapseExpandReactNode extends BaseTransform<CollapseExpandReactNo
     super(context, Object.assign({}, CollapseExpandReactNode.defaultOptions, options));
   }
 
-  public afterLayout() {
-    const { graph, element, model } = this.context;
+  public beforeDraw(input: DrawData): DrawData {
+    const { graph, element } = this.context;
     const { nodes = [], edges = [] } = graph.getData();
     const { enable, ...options } = this.options;
 
     nodes.forEach((datum) => {
       const nodeId = idOf(datum);
 
-      const node = element!.getElement(nodeId);
+      const node = element!.getElement(datum.id);
       if (!node || (datum.children && datum.children.length > 0)) return;
 
       const children = getNeighborNodeIds(nodeId, edges, this.options.direction);
       if (children.length === 0) return;
 
-      model.updateNodeData([{ id: nodeId, children }]);
+      datum.children = children;
     });
 
     const nodeMapper = graph.getOptions().node!;
@@ -114,6 +114,7 @@ export class CollapseExpandReactNode extends BaseTransform<CollapseExpandReactNo
     }
 
     graph.setNode(nodeMapper);
-    graph.draw();
+
+    return input;
   }
 }