feat(components): [tree-v2] add props.class prop (#18911)

* feat(components): [tree-v2] allow tree node to have customizable class

* docs: update docs

* docs: add demo

* docs: update

* fix: update

docs/en-US/component/tree-v2.md

@@ -57,6 +57,16 @@ tree-v2/custom-node
 
 :::
 
+## Custom node class ^(2.9.0)
+
+The class of tree nodes can be customized
+
+:::demo
+
+tree-v2/custom-node-class
+
+:::
+
 ## Tree node filtering
 
 Tree nodes can be filtered
@@ -90,11 +100,12 @@ tree-v2/filter
 ## props
 
 | Attribute      | Description                                                                          | Type                                             | Default  |
-| --------- | ------------------------------------------------------------------------------------ | -------------- | -------- |
-| value     | unique identity key name for nodes, its value should be unique across the whole tree | string, number | id       |
+| -------------- | ------------------------------------------------------------------------------------ | ------------------------------------------------ | -------- |
+| value          | unique identity key name for nodes, its value should be unique across the whole tree | string,number                                    | id       |
 | label          | specify which key of node object is used as the node's label                         | string                                           | label    |
 | children       | specify which node object is used as the node's subtree                              | string                                           | children |
 | disabled       | specify which key of node object represents if node's checkbox is disabled           | string                                           | disabled |
+| class ^(2.9.0) | custom node class name                                                               | ^[string] \| ^[Function]`(data, node) => string` | —        |
 
 ## TreeV2 Method
 

docs/examples/tree-v2/custom-node-class.vue

@@ -0,0 +1,87 @@
+<template>
+  <el-tree-v2
+    style="max-width: 600px"
+    :data="data"
+    show-checkbox
+    :expand-on-click-node="false"
+    :props="{ class: customNodeClass }"
+  />
+</template>
+
+<script lang="ts" setup>
+import type {
+  TreeNode,
+  TreeNodeData,
+} from 'element-plus/es/components/tree-v2/src/types'
+
+interface Tree {
+  id?: string
+  value?: string
+  label?: string
+  isPenultimate?: boolean
+  children?: Tree[]
+}
+
+const customNodeClass = ({ isPenultimate }: TreeNodeData, node: TreeNode) =>
+  isPenultimate ? 'is-penultimate' : ''
+
+const data: Tree[] = [
+  {
+    id: '1',
+    label: 'Level one 1',
+    children: [
+      {
+        id: '4',
+        label: 'Level two 1-1',
+        isPenultimate: true,
+        children: [
+          {
+            id: '9',
+            label: 'Level three 1-1-1',
+          },
+          {
+            id: '10',
+            label: 'Level three 1-1-2',
+          },
+        ],
+      },
+    ],
+  },
+  {
+    id: '2',
+    label: 'Level one 2',
+    isPenultimate: true,
+    children: [
+      {
+        id: '5',
+        label: 'Level two 2-1',
+      },
+      {
+        id: '6',
+        label: 'Level two 2-2',
+      },
+    ],
+  },
+  {
+    id: '3',
+    label: 'Level one 3',
+    isPenultimate: true,
+    children: [
+      {
+        id: '7',
+        label: 'Level two 3-1',
+      },
+      {
+        id: '8',
+        label: 'Level two 3-2',
+      },
+    ],
+  },
+]
+</script>
+
+<style>
+.is-penultimate > .el-tree-node__content {
+  color: var(--el-color-primary);
+}
+</style>

packages/components/tree-v2/__tests__/tree.test.ts

@@ -20,9 +20,7 @@ const TREE_NODE_CLASS_NAME = '.el-tree-node'
 const TREE_NODE_CONTENT_CLASS_NAME = '.el-tree-node__content'
 const TREE_NODE_EXPAND_ICON_CLASS_NAME = '.el-tree-node__expand-icon'
 
-const getUniqueId = () => {
-  return id++
-}
+const getUniqueId = () => id++
 
 const createData = (
   maxDeep,
@@ -823,6 +821,36 @@ describe('Virtual Tree', () => {
     expect(nodes[1].classes()).toContain('is-current')
   })
 
+  test('customNodeClass', async () => {
+    const { wrapper } = createTree({
+      data() {
+        return {
+          data: [
+            {
+              id: '1',
+              label: 'node-1',
+            },
+            {
+              id: '2',
+              label: 'node-2',
+            },
+          ],
+          props: {
+            value: 'id',
+            label: 'label',
+            children: 'children',
+            class: (data) => (data.id === '1' ? 'is-test' : ''),
+          },
+        }
+      },
+    })
+    await nextTick()
+    const currentNodeLabelWrapper = wrapper.find(
+      '.is-test .el-tree-node__label'
+    )
+    expect(currentNodeLabelWrapper.text()).toEqual('node-1')
+  })
+
   test('custom node content', async () => {
     const { wrapper } = createTree({
       slots: {

packages/components/tree-v2/src/tree-node.vue

@@ -7,6 +7,7 @@
       ns.is('current', current),
       ns.is('focusable', !disabled),
       ns.is('checked', !disabled && checked),
+      getNodeClass(node),
     ]"
     role="treeitem"
     tabindex="-1"
@@ -60,6 +61,7 @@ import ElIcon from '@element-plus/components/icon'
 import { CaretRight } from '@element-plus/icons-vue'
 import ElCheckbox from '@element-plus/components/checkbox'
 import { useNamespace } from '@element-plus/hooks'
+import { isFunction, isString } from '@element-plus/utils'
 import ElNodeContent from './tree-node-content'
 import {
   NODE_CONTEXTMENU,
@@ -68,6 +70,7 @@ import {
   treeNodeProps,
 } from './virtual-tree'
 import type { CheckboxValueType } from '@element-plus/components/checkbox'
+import type { TreeNode } from './types'
 
 defineOptions({
   name: 'ElTreeNode',
@@ -79,18 +82,27 @@ const emit = defineEmits(treeNodeEmits)
 const tree = inject(ROOT_TREE_INJECTION_KEY)
 const ns = useNamespace('tree')
 
-const indent = computed(() => {
-  return tree?.props.indent ?? 16
-})
+const indent = computed(() => tree?.props.indent ?? 16)
+const icon = computed(() => tree?.props.icon ?? CaretRight)
 
-const icon = computed(() => {
-  return tree?.props.icon ?? CaretRight
-})
+const getNodeClass = (node: TreeNode) => {
+  const nodeClassFunc = tree?.props.props.class
+  if (!nodeClassFunc) return {}
+
+  let className
+  if (isFunction(nodeClassFunc)) {
+    const { data } = node
+    className = nodeClassFunc(data, node)
+  } else {
+    className = nodeClassFunc
+  }
+
+  return isString(className) ? { [className]: true } : className
+}
 
 const handleClick = (e: MouseEvent) => {
   emit('click', props.node, e)
 }
-
 const handleDrop = (e: DragEvent) => {
   emit('drop', props.node, e)
 }
@@ -100,6 +112,7 @@ const handleExpandIconClick = () => {
 const handleCheckChange = (value: CheckboxValueType) => {
   emit('check', props.node, value)
 }
+
 const handleContextMenu = (event: Event) => {
   if (tree?.instance?.vnode?.props?.['onNodeContextmenu']) {
     event.stopPropagation()

packages/components/tree-v2/src/types.ts

@@ -16,6 +16,10 @@ export interface TreeOptionProps {
   label?: string
   value?: string
   disabled?: string
+  class?: (
+    data: TreeNodeData,
+    node: TreeNode
+  ) => string | { [key: string]: boolean }
 }
 
 export type TreeProps = ExtractPropTypes<typeof treeProps>

packages/components/tree-v2/src/virtual-tree.ts

@@ -32,6 +32,7 @@ export enum TreeOptionsEnum {
   LABEL = 'label',
   CHILDREN = 'children',
   DISABLED = 'disabled',
+  CLASS = '',
 }
 
 export const enum SetOperationEnum {
@@ -65,6 +66,7 @@ export const treeProps = buildProps({
         label: TreeOptionsEnum.LABEL,
         disabled: TreeOptionsEnum.DISABLED,
         value: TreeOptionsEnum.KEY,
+        class: TreeOptionsEnum.CLASS,
       } as const),
   },
   highlightCurrent: {