@ithinkdt/core/src/page/use-tree.js

iThinkDT Core

import { watch, computed, ref, shallowRef } from 'vue'

import { until } from '@vueuse/core'

import { array2Tree, walkTree } from '@ithinkdt/common'

export function useTree({ data, idKey = 'id', transform, defaultSelectRoot, defaultExpandRoot }) {
    data = shallowRef(data)
    const tree = computed(() => {
        if (transform === true) {
            transform = {}
        }
        if (transform) {
            transform.idKey = idKey
        }
        return transform ? array2Tree(data.value, transform) : data.value
    })

    const expandedKeys = ref([])
    const selectedKeys = ref([])
    const checkedKeys = ref([])

    until(tree)
        .toMatch((arr) => arr.length > 0)
        .then((data) => {
            if (defaultSelectRoot) {
                selectedKeys.value = data[0]?.[idKey] ? [data[0]?.[idKey]] : []
            }
        })

    watch(
        tree,
        (data) => {
            if (defaultExpandRoot) {
                expandedKeys.value = data.map((it) => it?.[idKey])
            }
        },
        { immediate: true },
    )

    return {
        tree,
        selectedKeys,
        checkedKeys,
        expandedKeys,
        getByKey(key) {
            let it
            walkTree(tree.value, (it2) => {
                if (it2[idKey] === key) {
                    it = it2
                    return false
                }
            })

            return it
        },
        getParentByKey(key) {
            let it
            walkTree(tree.value, (it2, _i, parent) => {
                if (it2[idKey] === key) {
                    it = parent
                    return false
                }
            })

            return it
        },
    }
}