53 lines
1.5 KiB
53 lines
1.5 KiB
import dagre from 'dagre';
|
|
const edgeType = 'smoothstep';
|
|
|
|
export const flattenTree = (tree) => {
|
|
let stack = [tree];
|
|
let nodes = [];
|
|
let edges = [];
|
|
while (stack.length > 0) {
|
|
let node = stack.pop();
|
|
nodes.push({ id: node.id, data: { ...node, label: `${node.name}` }, position: node.position });
|
|
if (node.children && node.children.length > 0) {
|
|
node.children.map(item => {
|
|
edges.push({ source: node.id, target: item.id, id: item.id, type: edgeType, animated: true });
|
|
stack.unshift(item);
|
|
});
|
|
};
|
|
};
|
|
return { nodes, edges };
|
|
};
|
|
|
|
export const getLayoutedElements = (nodes, edges, direction = 'TB') => {
|
|
const dagreGraph = new dagre.graphlib.Graph();
|
|
dagreGraph.setDefaultEdgeLabel(() => ({}));
|
|
|
|
const nodeWidth = 172;
|
|
const nodeHeight = 36;
|
|
const isHorizontal = direction === 'LR';
|
|
dagreGraph.setGraph({ rankdir: direction });
|
|
|
|
nodes.forEach((node: any) => {
|
|
dagreGraph.setNode(node.id, { width: nodeWidth, height: nodeHeight });
|
|
});
|
|
|
|
edges.forEach((edge: any) => {
|
|
dagreGraph.setEdge(edge.source, edge.target);
|
|
});
|
|
|
|
dagre.layout(dagreGraph);
|
|
|
|
nodes.forEach((node: any) => {
|
|
const nodeWithPosition = dagreGraph.node(node.id);
|
|
node.targetPosition = isHorizontal ? 'left' : 'top';
|
|
node.sourcePosition = isHorizontal ? 'right' : 'bottom';
|
|
node.position = {
|
|
x: nodeWithPosition.x - nodeWidth / 2,
|
|
y: nodeWithPosition.y - nodeHeight / 2,
|
|
};
|
|
|
|
return node;
|
|
});
|
|
|
|
return { nodes, edges };
|
|
};
|