/* * @Description: * @Author: leizhe * @Date: 2022-03-04 18:33:24 * @LastEditTime: 2022-04-27 14:26:14 * @LastEditors: leizhe */ import { Tree } from 'antd'; import { uniq, remove } from 'lodash'; import React, { useState, useEffect } from 'react'; const getChildKeys = tree => { let childKeys = tree.children?.map(c => getChildKeys(c)).flat(Infinity) || []; return [tree.key, ...childKeys]; }; export default props => { const { expandedKeys, onExpand, onSelect, selectedKeys, keepTree, ...rest } = props; const [epKeys, setEpKeys] = useState([expandedKeys]); useEffect(() => { console.log(expandedKeys); console.log(selectedKeys); setEpKeys([expandedKeys]); }, [expandedKeys]); const handleExpand = (keys, { expanded, node }) => { if (keepTree && keepTree.indexOf(node.key) != -1 && epKeys.indexOf(node.key) == -1) { setEpKeys([node.key]); } else { if (!expanded) { let childKeys = getChildKeys(node); remove(keys, k => childKeys.includes(k)); } setEpKeys(keys); } }; const handleSelect = (keys, e) => { if (keepTree && keepTree.indexOf(keys[0]) != -1 && epKeys.indexOf(keys[0]) == -1) { setEpKeys(keys); } else { let result = []; if (keys.length > 0 && epKeys.includes(keys[0])) { remove(epKeys, k => keys.includes(k)); result = [...epKeys]; } else if (keys[0]) { result = [...epKeys, keys[0]]; } else { if (epKeys.includes(e.node.key)) { remove(epKeys, k => k === e.node.key); } else { epKeys.push(e.node.key); } result = [...epKeys]; } console.log(result, 'result'); setEpKeys(result); } // eslint-disable-next-line no-unused-expressions onSelect && onSelect(keys, e); }; return ( <Tree onSelect={handleSelect} onExpand={handleExpand} expandedKeys={epKeys} selectedKeys={selectedKeys} {...rest} /> ); };