From b61236572bbba1d4dd592c8b0f08dfbb2891bf55 Mon Sep 17 00:00:00 2001
From: yvonneyx <zhuyuxin0627@gmail.com>
Date: Mon, 23 Dec 2024 15:37:40 +0800
Subject: [PATCH] refactor: collapse-expand-react-node transform

---
 .../core/transform/collapse-expand-react-node.tsx   | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/packages/graphs/src/core/transform/collapse-expand-react-node.tsx b/packages/graphs/src/core/transform/collapse-expand-react-node.tsx
index 3339c4658..cbe02251a 100644
--- a/packages/graphs/src/core/transform/collapse-expand-react-node.tsx
+++ b/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;
   }
 }