From 60e5db1e498023b0ceeb6600c8c138218fef517e Mon Sep 17 00:00:00 2001
From: snowingfox <1503401882@qq.com>
Date: Fri, 21 Oct 2022 09:47:18 +0800
Subject: [PATCH] refactor(components): [tree-v2] switch to script-setup syntax
 (#10133)

---
 .../tree-v2/src/composables/useCheck.ts       |   7 +-
 .../tree-v2/src/composables/useTree.ts        |  13 +-
 packages/components/tree-v2/src/tree-node.vue |  87 ++++------
 packages/components/tree-v2/src/tree.vue      | 151 ++++++++----------
 packages/components/tree-v2/src/types.ts      |   2 +-
 .../components/tree-v2/src/virtual-tree.ts    |   5 +-
 6 files changed, 120 insertions(+), 145 deletions(-)

diff --git a/packages/components/tree-v2/src/composables/useCheck.ts b/packages/components/tree-v2/src/composables/useCheck.ts
index 33be1c89e1..d462d18d6f 100644
--- a/packages/components/tree-v2/src/composables/useCheck.ts
+++ b/packages/components/tree-v2/src/composables/useCheck.ts
@@ -4,6 +4,7 @@ import {
   NODE_CHECK_CHANGE,
   SetOperationEnum,
 } from '../virtual-tree'
+import type { CheckboxValueType } from '@element-plus/components/checkbox'
 import type { Ref } from 'vue'
 import type { Tree, TreeKey, TreeNode, TreeNodeData, TreeProps } from '../types'
 
@@ -78,11 +79,11 @@ export function useCheck(props: TreeProps, tree: Ref<Tree | undefined>) {
 
   const toggleCheckbox = (
     node: TreeNode,
-    isChecked: boolean,
+    isChecked: CheckboxValueType,
     nodeClick = true
   ) => {
     const checkedKeySet = checkedKeys.value
-    const toggle = (node: TreeNode, checked: boolean) => {
+    const toggle = (node: TreeNode, checked: CheckboxValueType) => {
       checkedKeySet[checked ? SetOperationEnum.ADD : SetOperationEnum.DELETE](
         node.key
       )
@@ -102,7 +103,7 @@ export function useCheck(props: TreeProps, tree: Ref<Tree | undefined>) {
     }
   }
 
-  const afterNodeCheck = (node: TreeNode, checked: boolean) => {
+  const afterNodeCheck = (node: TreeNode, checked: CheckboxValueType) => {
     const { checkedNodes, checkedKeys } = getChecked()
     const { halfCheckedNodes, halfCheckedKeys } = getHalfChecked()
     emit(NODE_CHECK, node.data, {
diff --git a/packages/components/tree-v2/src/composables/useTree.ts b/packages/components/tree-v2/src/composables/useTree.ts
index a8f9194490..f3d109bb41 100644
--- a/packages/components/tree-v2/src/composables/useTree.ts
+++ b/packages/components/tree-v2/src/composables/useTree.ts
@@ -1,4 +1,3 @@
-// @ts-nocheck
 import { computed, nextTick, ref, shallowRef, watch } from 'vue'
 import { isObject } from '@element-plus/utils'
 import {
@@ -10,6 +9,9 @@ import {
 } from '../virtual-tree'
 import { useCheck } from './useCheck'
 import { useFilter } from './useFilter'
+import type { SetupContext } from 'vue'
+import type { treeEmits } from '../virtual-tree'
+import type { CheckboxValueType } from '@element-plus/components/checkbox'
 import type {
   Tree,
   TreeData,
@@ -19,7 +21,10 @@ import type {
   TreeProps,
 } from '../types'
 
-export function useTree(props: TreeProps, emit) {
+export function useTree(
+  props: TreeProps,
+  emit: SetupContext<typeof treeEmits>['emit']
+) {
   const expandedKeySet = ref<Set<TreeKey>>(new Set(props.defaultExpandedKeys))
   const currentKey = ref<TreeKey | undefined>()
   const tree = shallowRef<Tree | undefined>()
@@ -213,7 +218,7 @@ export function useTree(props: TreeProps, emit) {
     }
   }
 
-  function handleNodeCheck(node: TreeNode, checked: boolean) {
+  function handleNodeCheck(node: TreeNode, checked: CheckboxValueType) {
     toggleCheckbox(node, checked)
   }
 
@@ -224,7 +229,7 @@ export function useTree(props: TreeProps, emit) {
       const { treeNodeMap } = tree.value
       keySet.forEach((key) => {
         const treeNode = treeNodeMap.get(key)
-        if (node && node.level === treeNode.level) {
+        if (node && node.level === treeNode?.level) {
           keySet.delete(key)
         }
       })
diff --git a/packages/components/tree-v2/src/tree-node.vue b/packages/components/tree-v2/src/tree-node.vue
index d001ac39ca..10aed03803 100644
--- a/packages/components/tree-v2/src/tree-node.vue
+++ b/packages/components/tree-v2/src/tree-node.vue
@@ -48,10 +48,8 @@
   </div>
 </template>
 
-<script lang="ts">
-// @ts-nocheck
-import { computed, defineComponent, inject } from 'vue'
-import { CaretRight } from '@element-plus/icons-vue'
+<script lang="ts" setup>
+import { computed, inject } from 'vue'
 import ElIcon from '@element-plus/components/icon'
 import ElCheckbox from '@element-plus/components/checkbox'
 import { useNamespace } from '@element-plus/hooks'
@@ -62,57 +60,42 @@ import {
   treeNodeEmits,
   treeNodeProps,
 } from './virtual-tree'
+import type { CheckboxValueType } from '@element-plus/components/checkbox'
 
 const DEFAULT_ICON = 'caret-right'
 
-export default defineComponent({
+defineOptions({
   name: 'ElTreeNode',
-  components: {
-    ElIcon,
-    CaretRight,
-    ElCheckbox,
-    ElNodeContent,
-  },
-  props: treeNodeProps,
-  emits: treeNodeEmits,
-  setup(props, { emit }) {
-    const tree = inject(ROOT_TREE_INJECTION_KEY)
-    const ns = useNamespace('tree')
-
-    const indent = computed(() => {
-      return tree?.props.indent ?? 16
-    })
-
-    const icon = computed(() => {
-      return tree?.props.icon ?? DEFAULT_ICON
-    })
-
-    const handleClick = (e: MouseEvent) => {
-      emit('click', props.node, e)
-    }
-    const handleExpandIconClick = () => {
-      emit('toggle', props.node)
-    }
-    const handleCheckChange = (value: boolean) => {
-      emit('check', props.node, value)
-    }
-    const handleContextMenu = (event: Event) => {
-      if (tree?.instance?.vnode?.props?.['onNodeContextmenu']) {
-        event.stopPropagation()
-        event.preventDefault()
-      }
-      tree?.ctx.emit(NODE_CONTEXTMENU, event, props.node?.data, props.node)
-    }
-
-    return {
-      ns,
-      indent,
-      icon,
-      handleClick,
-      handleExpandIconClick,
-      handleCheckChange,
-      handleContextMenu,
-    }
-  },
 })
+
+const props = defineProps(treeNodeProps)
+const emit = defineEmits(treeNodeEmits)
+
+const tree = inject(ROOT_TREE_INJECTION_KEY)
+const ns = useNamespace('tree')
+
+const indent = computed(() => {
+  return tree?.props.indent ?? 16
+})
+
+const icon = computed(() => {
+  return tree?.props.icon ?? DEFAULT_ICON
+})
+
+const handleClick = (e: MouseEvent) => {
+  emit('click', props.node, e)
+}
+const handleExpandIconClick = () => {
+  emit('toggle', props.node)
+}
+const handleCheckChange = (value: CheckboxValueType) => {
+  emit('check', props.node, value)
+}
+const handleContextMenu = (event: Event) => {
+  if (tree?.instance?.vnode?.props?.['onNodeContextmenu']) {
+    event.stopPropagation()
+    event.preventDefault()
+  }
+  tree?.ctx.emit(NODE_CONTEXTMENU, event, props.node?.data, props.node)
+}
 </script>
diff --git a/packages/components/tree-v2/src/tree.vue b/packages/components/tree-v2/src/tree.vue
index f26ebf8743..3c61b041fc 100644
--- a/packages/components/tree-v2/src/tree.vue
+++ b/packages/components/tree-v2/src/tree.vue
@@ -38,99 +38,84 @@
   </div>
 </template>
 
-<script lang="ts">
-import { defineComponent, getCurrentInstance, provide } from 'vue'
+<script lang="ts" setup>
+import { getCurrentInstance, provide, useSlots } from 'vue'
 import { useLocale, useNamespace } from '@element-plus/hooks'
 import { formItemContextKey } from '@element-plus/tokens'
 import { FixedSizeList } from '@element-plus/components/virtual-list'
 import { useTree } from './composables/useTree'
 import ElTreeNode from './tree-node.vue'
 import { ROOT_TREE_INJECTION_KEY, treeEmits, treeProps } from './virtual-tree'
-import type { TreeProps } from './types'
 
-export default defineComponent({
+defineOptions({
   name: 'ElTreeV2',
-  components: {
-    ElTreeNode,
-    FixedSizeList,
-  },
-  props: treeProps,
-  emits: treeEmits,
-  setup(props: TreeProps, ctx) {
-    provide(ROOT_TREE_INJECTION_KEY, {
-      ctx,
-      props,
-      instance: getCurrentInstance(),
-    })
-    provide(formItemContextKey, undefined)
-    const { t } = useLocale()
-    const ns = useNamespace('tree')
-    const {
-      flattenTree,
-      isNotEmpty,
-      toggleExpand,
-      isExpanded,
-      isIndeterminate,
-      isChecked,
-      isDisabled,
-      isCurrent,
-      isForceHiddenExpandIcon,
-      toggleCheckbox,
-      handleNodeClick,
-      handleNodeCheck,
-      // expose
-      getCurrentNode,
-      getCurrentKey,
-      setCurrentKey,
-      getCheckedKeys,
-      getCheckedNodes,
-      getHalfCheckedKeys,
-      getHalfCheckedNodes,
-      setChecked,
-      setCheckedKeys,
-      filter,
-      setData,
-      getNode,
-      expandNode,
-      collapseNode,
-      setExpandedKeys,
-    } = useTree(props, ctx.emit)
+})
 
-    ctx.expose({
-      getCurrentNode,
-      getCurrentKey,
-      setCurrentKey,
-      getCheckedKeys,
-      getCheckedNodes,
-      getHalfCheckedKeys,
-      getHalfCheckedNodes,
-      setChecked,
-      setCheckedKeys,
-      filter,
-      setData,
-      getNode,
-      expandNode,
-      collapseNode,
-      setExpandedKeys,
-    })
+const props = defineProps(treeProps)
+const emit = defineEmits(treeEmits)
 
-    return {
-      t,
-      ns,
-      flattenTree,
-      itemSize: 26,
-      isNotEmpty,
-      toggleExpand,
-      toggleCheckbox,
-      isExpanded,
-      isIndeterminate,
-      isChecked,
-      isDisabled,
-      isCurrent,
-      isForceHiddenExpandIcon,
-      handleNodeClick,
-      handleNodeCheck,
-    }
+const slots = useSlots()
+
+const itemSize = 26
+
+provide(ROOT_TREE_INJECTION_KEY, {
+  ctx: {
+    emit,
+    slots,
   },
+  props,
+  instance: getCurrentInstance()!,
+})
+provide(formItemContextKey, undefined)
+const { t } = useLocale()
+const ns = useNamespace('tree')
+const {
+  flattenTree,
+  isNotEmpty,
+  toggleExpand,
+  isExpanded,
+  isIndeterminate,
+  isChecked,
+  isDisabled,
+  isCurrent,
+  isForceHiddenExpandIcon,
+  handleNodeClick,
+  handleNodeCheck,
+  // expose
+  toggleCheckbox,
+  getCurrentNode,
+  getCurrentKey,
+  setCurrentKey,
+  getCheckedKeys,
+  getCheckedNodes,
+  getHalfCheckedKeys,
+  getHalfCheckedNodes,
+  setChecked,
+  setCheckedKeys,
+  filter,
+  setData,
+  getNode,
+  expandNode,
+  collapseNode,
+  setExpandedKeys,
+} = useTree(props, emit)
+
+defineExpose({
+  toggleCheckbox,
+  getCurrentNode,
+  getCurrentKey,
+  setCurrentKey,
+  getCheckedKeys,
+  getCheckedNodes,
+  getHalfCheckedKeys,
+  getHalfCheckedNodes,
+  setChecked,
+  setCheckedKeys,
+  filter,
+  setData,
+  getNode,
+  expandNode,
+  collapseNode,
+  setExpandedKeys,
 })
 </script>
diff --git a/packages/components/tree-v2/src/types.ts b/packages/components/tree-v2/src/types.ts
index e692e72bc1..b1a9fc7401 100644
--- a/packages/components/tree-v2/src/types.ts
+++ b/packages/components/tree-v2/src/types.ts
@@ -32,7 +32,7 @@ export interface TreeNode {
 }
 
 export interface TreeContext {
-  ctx: SetupContext<typeof treeEmits>
+  ctx: Omit<SetupContext<typeof treeEmits>, 'expose' | 'attrs'>
   instance: ComponentInternalInstance
   props: TreeProps
 }
diff --git a/packages/components/tree-v2/src/virtual-tree.ts b/packages/components/tree-v2/src/virtual-tree.ts
index 330a160308..18d1194e0c 100644
--- a/packages/components/tree-v2/src/virtual-tree.ts
+++ b/packages/components/tree-v2/src/virtual-tree.ts
@@ -4,8 +4,9 @@ import {
   iconPropType,
   mutable,
 } from '@element-plus/utils'
+import type { CheckboxValueType } from '@element-plus/components/checkbox'
 import type { InjectionKey } from 'vue'
-import type { TreeNodeData } from '../../tree/src/tree.type'
+import type { TreeNodeData } from '@element-plus/components/tree/src/tree.type'
 import type {
   CheckedInfo,
   FilterMethod,
@@ -183,6 +184,6 @@ export const treeEmits = {
 export const treeNodeEmits = {
   click: (node: TreeNode, e: MouseEvent) => !!(node && e),
   toggle: (node: TreeNode) => !!node,
-  check: (node: TreeNode, checked: boolean) =>
+  check: (node: TreeNode, checked: CheckboxValueType) =>
     node && typeof checked === 'boolean',
 }