index.jsx

/*
 * @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);
    if (!expandedKeys || expandedKeys == null) {
      setEpKeys(['0-0']);
    } else {
      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}
    />
  );
};