From ccfa7aa2699f1715a006e7ee6b1504aa92f07226 Mon Sep 17 00:00:00 2001
From: Lukas Bach <contact@lukasbach.com>
Date: Sat, 2 Mar 2024 20:22:08 +0100
Subject: [PATCH] feat: shouldRenderChildren prop (#333)

---
 next-release-notes.md                         |  4 +-
 .../src/stories/BasicExamples.stories.tsx     | 64 +++++++++++++++++++
 .../core/src/treeItem/TreeItemElement.tsx     |  6 +-
 .../src/treeItem/useTreeItemRenderContext.ts  |  9 ++-
 packages/core/src/types.ts                    |  4 ++
 5 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/next-release-notes.md b/next-release-notes.md
index b90e66d63..9d4befe49 100644
--- a/next-release-notes.md
+++ b/next-release-notes.md
@@ -2,7 +2,9 @@
 
 - Add `setDomFocus` argument to focus-item methods to provide an escape hatch to set the focus state of an item in RCT
   without updating the DOM focus. This defaults to true in all existing methods to maintain the current behavior if
-  it is absent.
+  it is absent. (#336)
+- Allow customizing when a subtree is rendered or not with the new `shouldRenderChildren` prop. This can be used to 
+  create opening and closing animations on subtrees. (#333)
 
 ### Bug Fixes
 
diff --git a/packages/core/src/stories/BasicExamples.stories.tsx b/packages/core/src/stories/BasicExamples.stories.tsx
index ded604283..02f9e711b 100644
--- a/packages/core/src/stories/BasicExamples.stories.tsx
+++ b/packages/core/src/stories/BasicExamples.stories.tsx
@@ -7,6 +7,7 @@ import { Tree } from '../tree/Tree';
 import { StaticTreeDataProvider } from '../uncontrolledEnvironment/StaticTreeDataProvider';
 import { UncontrolledTreeEnvironment } from '../uncontrolledEnvironment/UncontrolledTreeEnvironment';
 import { buildTestTree } from '../../test/helpers';
+import { TreeItemIndex } from '../types';
 
 export default {
   title: 'Core/Basic Examples',
@@ -556,3 +557,66 @@ export const NavigateAway = () => {
     </UncontrolledTreeEnvironment>
   );
 };
+
+export const AnimatedExpandingAndCollapsing = () => {
+  const [openingItem, setOpeningItem] = useState<TreeItemIndex | undefined>();
+  const [closingItem, setClosingItem] = useState<TreeItemIndex | undefined>();
+  return (
+    <UncontrolledTreeEnvironment<string>
+      dataProvider={
+        new StaticTreeDataProvider(longTree.items, (item, data) => ({
+          ...item,
+          data,
+        }))
+      }
+      getItemTitle={item => item.data}
+      shouldRenderChildren={(item, context) =>
+        context.isExpanded ||
+        closingItem === item.index ||
+        openingItem === item.index
+      }
+      onExpandItem={item => {
+        setOpeningItem(item.index);
+        setTimeout(() => {
+          setOpeningItem(undefined);
+        });
+      }}
+      onCollapseItem={item => {
+        setClosingItem(item.index);
+        setTimeout(() => {
+          setClosingItem(undefined);
+        }, 500);
+      }}
+      viewState={{
+        'tree-1': {},
+      }}
+      renderItemsContainer={({ children, containerProps, parentId }) => (
+        <div
+          style={{
+            overflow: 'hidden',
+          }}
+        >
+          <ul
+            {...containerProps}
+            className="rct-tree-items-container"
+            style={{
+              transition: 'all 500ms',
+              maxHeight:
+                parentId === openingItem || parentId === closingItem
+                  ? 0
+                  : '999999px',
+              marginTop:
+                parentId === closingItem || parentId === openingItem
+                  ? '-100%'
+                  : '0',
+            }}
+          >
+            {children}
+          </ul>
+        </div>
+      )}
+    >
+      <Tree treeId="tree-1" rootItem="root" treeLabel="Tree Example" />
+    </UncontrolledTreeEnvironment>
+  );
+};
diff --git a/packages/core/src/treeItem/TreeItemElement.tsx b/packages/core/src/treeItem/TreeItemElement.tsx
index f006a1a4f..00fafda3d 100644
--- a/packages/core/src/treeItem/TreeItemElement.tsx
+++ b/packages/core/src/treeItem/TreeItemElement.tsx
@@ -32,7 +32,11 @@ export const TreeItemElement = (props: {
     return null as any;
   }
 
-  const children = item.isFolder && isExpanded && item.children && (
+  const shouldRenderChildren =
+    environment.shouldRenderChildren?.(item, renderContext) ??
+    (item.isFolder && isExpanded);
+
+  const children = item.children && shouldRenderChildren && (
     <TreeItemChildren depth={props.depth + 1} parentId={props.itemIndex}>
       {item.children}
     </TreeItemChildren>
diff --git a/packages/core/src/treeItem/useTreeItemRenderContext.ts b/packages/core/src/treeItem/useTreeItemRenderContext.ts
index 97675abcb..dbd94d0e1 100644
--- a/packages/core/src/treeItem/useTreeItemRenderContext.ts
+++ b/packages/core/src/treeItem/useTreeItemRenderContext.ts
@@ -1,5 +1,10 @@
 import { HTMLProps, useMemo } from 'react';
-import { TreeItem, TreeItemActions, TreeItemRenderFlags } from '../types';
+import {
+  TreeItem,
+  TreeItemActions,
+  TreeItemRenderContext,
+  TreeItemRenderFlags,
+} from '../types';
 import { defaultMatcher } from '../search/defaultMatcher';
 import { useTree } from '../tree/Tree';
 import { useTreeEnvironment } from '../controlledEnvironment/ControlledTreeEnvironment';
@@ -37,7 +42,7 @@ export const useTreeItemRenderContext = (item?: TreeItem) => {
     item && environment.viewState[treeId]?.expandedItems?.includes(item.index);
   const isRenaming = item && renamingItem === item.index;
 
-  return useMemo(() => {
+  return useMemo<TreeItemRenderContext | undefined>(() => {
     if (!item) {
       return undefined;
     }
diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts
index c14e4b519..77fe15305 100644
--- a/packages/core/src/types.ts
+++ b/packages/core/src/types.ts
@@ -178,6 +178,10 @@ export interface TreeCapabilities<T = any, C extends string = never> {
     itemTitle: string
   ) => boolean;
   showLiveDescription?: boolean;
+  shouldRenderChildren?: (
+    item: TreeItem<T>,
+    context: TreeItemRenderContext<C>
+  ) => boolean;
 
   /**
    * See Issue #148 or the sample at