feat: add disabledStrategy prop to support dynamic node disabled state calculation

src/TreeNode.tsx

@@ -2,7 +2,7 @@ import classNames from 'classnames';
 import pickAttrs from 'rc-util/lib/pickAttrs';
 import * as React from 'react';
 // @ts-ignore
-import { TreeContext, type TreeContextProps } from './contextTypes';
+import { TreeContext, UnstableContext, type TreeContextProps } from './contextTypes';
 import Indent from './Indent';
 import type { TreeNodeProps } from './interface';
 import getEntity from './utils/keyUtil';
@@ -17,6 +17,7 @@ export type { TreeNodeProps } from './interface';
 
 export interface InternalTreeNodeProps extends TreeNodeProps {
   context?: TreeContextProps;
+  unstableContext?: React.ContextType<typeof UnstableContext>;
 }
 
 export interface TreeNodeState {
@@ -231,12 +232,13 @@ class InternalTreeNode extends React.Component<InternalTreeNodeProps, TreeNodeSt
   };
 
   isDisabled = () => {
-    const { disabled } = this.props;
+    const { disabled, data } = this.props;
     const {
       context: { disabled: treeDisabled },
+      unstableContext: { nodeDisabled },
     } = this.props;
 
-    return !!(treeDisabled || disabled);
+    return !!(treeDisabled || disabled || (nodeDisabled && nodeDisabled(data)));
   };
 
   isCheckable = () => {
@@ -596,7 +598,13 @@ class InternalTreeNode extends React.Component<InternalTreeNodeProps, TreeNodeSt
 
 const ContextTreeNode: React.FC<TreeNodeProps> = props => (
   <TreeContext.Consumer>
-    {context => <InternalTreeNode {...props} context={context} />}
+    {context => (
+      <UnstableContext.Consumer>
+        {unstableContext => (
+          <InternalTreeNode {...props} context={context} unstableContext={unstableContext} />
+        )}
+      </UnstableContext.Consumer>
+    )}
   </TreeContext.Consumer>
 );
 

src/contextTypes.ts

@@ -93,3 +93,8 @@ export interface TreeContextProps<TreeDataType extends BasicDataNode = DataNode>
 }
 
 export const TreeContext: React.Context<TreeContextProps<any> | null> = React.createContext(null);
+
+/** Internal usage, safe to remove. Do not use in prod */
+export const UnstableContext = React.createContext<{
+  nodeDisabled?: (node: DataNode) => boolean;
+}>({});

src/index.ts

@@ -2,7 +2,8 @@ import Tree from './Tree';
 import TreeNode from './TreeNode';
 import type { TreeProps } from './Tree';
 import type { TreeNodeProps, BasicDataNode, FieldDataNode } from './interface';
+import { UnstableContext } from './contextTypes';
 
-export { TreeNode };
+export { TreeNode, UnstableContext };
 export type { TreeProps, TreeNodeProps, BasicDataNode, FieldDataNode };
 export default Tree;

src/interface.tsx

@@ -69,7 +69,10 @@ export type Key = React.Key;
  */
 export type SafeKey = Exclude<Key, bigint>;
 
-export type KeyEntities<DateType extends BasicDataNode = any> = Record<SafeKey, DataEntity<DateType>>;
+export type KeyEntities<DateType extends BasicDataNode = any> = Record<
+  SafeKey,
+  DataEntity<DateType>
+>;
 
 export type DataNode = FieldDataNode<{
   key: Key;

tests/Tree.spec.tsx

@@ -1,11 +1,12 @@
 /* eslint-disable no-undef, react/no-multi-comp, no-console,
 react/no-unused-state, react/prop-types, no-return-assign */
 import React from 'react';
-import {render, fireEvent, act} from '@testing-library/react';
+import { render, fireEvent, act } from '@testing-library/react';
 import { resetWarned } from 'rc-util/lib/warning';
 import { spyElementPrototypes } from 'rc-util/lib/test/domHook';
 import Tree, { TreeNode } from '../src';
 import { objectMatcher, spyConsole, spyError } from './util';
+import { UnstableContext } from '../src';
 
 const OPEN_CLASSNAME = 'rc-tree-switcher_open';
 const CHECKED_CLASSNAME = 'rc-tree-checkbox-checked';
@@ -1055,9 +1056,9 @@ describe('Tree Basic', () => {
       render(<Tree ref={treeRef} />);
 
       act(() => {
-        treeRef.current.scrollTo({key: 'light', align: 'top'});
+        treeRef.current.scrollTo({ key: 'light', align: 'top' });
       });
-      
+
       jest.runAllTimers();
 
       expect(called).toBeTruthy();
@@ -1092,7 +1093,7 @@ describe('Tree Basic', () => {
       }
 
       const treeRef = React.createRef<any>();
-      
+
       const { container } = render(
         <Tree ref={treeRef} treeData={data} virtual itemHeight={20} height={100} />,
       );
@@ -1101,7 +1102,7 @@ describe('Tree Basic', () => {
         treeRef.current.scrollTo({ key: 5, offset: 5 });
         jest.runAllTimers();
       });
-      
+
       expect(container.querySelector('.rc-tree-list-holder').scrollTop).toEqual(25);
 
       jest.useRealTimers();
@@ -1237,4 +1238,96 @@ describe('Tree Basic', () => {
     expect(container.firstChild).toHaveClass('root-tree');
     expect(container.firstChild).toHaveStyle({ backgroundColor: 'cyan' });
   });
+
+  describe('nodeDisabled', () => {
+    it('should disable node by nodeDisabled function', () => {
+      const { getByRole } = render(
+        <UnstableContext.Provider
+          value={{
+            nodeDisabled: node => node.key === '0-0-0',
+          }}
+        >
+          <Tree defaultExpandAll>
+            <TreeNode title="parent 1" key="0-0">
+              <TreeNode title="leaf 1" key="0-0-0" />
+              <TreeNode title="leaf 2" key="0-0-1" />
+            </TreeNode>
+          </Tree>
+        </UnstableContext.Provider>,
+      );
+
+      expect(getByRole('treeitem', { name: 'leaf 1' })).toHaveClass('rc-tree-treenode-disabled');
+      expect(getByRole('treeitem', { name: 'leaf 2' })).not.toHaveClass(
+        'rc-tree-treenode-disabled',
+      );
+    });
+
+    it('should work with checkable tree', () => {
+      const onCheck = jest.fn();
+      const { container } = render(
+        <UnstableContext.Provider
+          value={{
+            nodeDisabled: node => node.key === '0-0-0',
+          }}
+        >
+          <Tree checkable defaultExpandAll onCheck={onCheck}>
+            <TreeNode title="parent 1" key="0-0">
+              <TreeNode title="leaf 1" key="0-0-0" />
+              <TreeNode title="leaf 2" key="0-0-1" />
+            </TreeNode>
+          </Tree>
+        </UnstableContext.Provider>,
+      );
+
+      fireEvent.click(container.querySelectorAll('.rc-tree-checkbox')[1]);
+      expect(onCheck).not.toHaveBeenCalled();
+
+      fireEvent.click(container.querySelectorAll('.rc-tree-checkbox')[2]);
+      expect(onCheck).toHaveBeenCalled();
+    });
+
+    it('should work with selectable tree', () => {
+      const onSelect = jest.fn();
+      const { container } = render(
+        <UnstableContext.Provider
+          value={{
+            nodeDisabled: node => node.key === '0-0-0',
+          }}
+        >
+          <Tree selectable defaultExpandAll onSelect={onSelect}>
+            <TreeNode title="parent 1" key="0-0">
+              <TreeNode title="leaf 1" key="0-0-0" />
+              <TreeNode title="leaf 2" key="0-0-1" />
+            </TreeNode>
+          </Tree>
+        </UnstableContext.Provider>,
+      );
+
+      fireEvent.click(container.querySelectorAll('.rc-tree-node-content-wrapper')[1]);
+      expect(onSelect).not.toHaveBeenCalled();
+
+      fireEvent.click(container.querySelectorAll('.rc-tree-node-content-wrapper')[2]);
+      expect(onSelect).toHaveBeenCalled();
+    });
+
+    it('should override disabled prop', () => {
+      const { getByRole } = render(
+        <UnstableContext.Provider
+          value={{
+            nodeDisabled: node => node.key === '0-0-0',
+          }}
+        >
+          <Tree defaultExpandAll>
+            <TreeNode title="parent 1" key="0-0">
+              <TreeNode title="leaf 1" key="0-0-0" disabled={false} />
+              <TreeNode title="leaf 2" key="0-0-1" disabled />
+            </TreeNode>
+          </Tree>
+        </UnstableContext.Provider>,
+      );
+
+      expect(getByRole('treeitem', { name: 'leaf 1' })).toHaveClass('rc-tree-treenode-disabled');
+      expect(getByRole('treeitem', { name: 'leaf 2' })).toHaveClass('rc-tree-treenode-disabled');
+    });
+  });
 });