From 28722416ee66ee13a039696ba84a1e241607c5ca Mon Sep 17 00:00:00 2001
From: Aarya B C <60278134+ayy-bc@users.noreply.github.com>
Date: Wed, 8 May 2024 19:39:22 +0000
Subject: [PATCH] change LeadingAction of type React.FC<TreeViewVisualProps>

---
 .../TreeView/TreeView.examples.stories.tsx    | 23 +++++++++++--------
 .../TreeView/TreeView.features.stories.tsx    | 12 +++++++---
 packages/react/src/TreeView/TreeView.tsx      | 19 +++++++++------
 3 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/packages/react/src/TreeView/TreeView.examples.stories.tsx b/packages/react/src/TreeView/TreeView.examples.stories.tsx
index 56d5130c1ed..5d047f306aa 100644
--- a/packages/react/src/TreeView/TreeView.examples.stories.tsx
+++ b/packages/react/src/TreeView/TreeView.examples.stories.tsx
@@ -3,6 +3,7 @@ import type {Meta, Story} from '@storybook/react'
 import React from 'react'
 import Box from '../Box'
 import {TreeView} from './TreeView'
+import {IconButton} from '../Button'
 
 const meta: Meta = {
   title: 'Components/TreeView/Examples',
@@ -53,16 +54,18 @@ const ControlledDraggableItem: React.FC<{id: string; children: React.ReactNode}>
   return (
     <>
       <TreeView.Item id={id} className="treeview-item" expanded={expanded} onExpandedChange={setExpanded}>
-        <TreeView.LeadingAction
-          icon={GrabberIcon}
-          aria-label="Reorder item"
-          className="treeview-leading-action"
-          draggable="true"
-          onDragStart={() => {
-            setExpanded(false)
-            // other drag logic to follow
-          }}
-        />
+        <TreeView.LeadingAction>
+          <IconButton
+            icon={GrabberIcon}
+            aria-label="Reorder item"
+            className="treeview-leading-action"
+            draggable="true"
+            onDragStart={() => {
+              setExpanded(false)
+              // other drag logic to follow
+            }}
+          />
+        </TreeView.LeadingAction>
         {children}
       </TreeView.Item>
     </>
diff --git a/packages/react/src/TreeView/TreeView.features.stories.tsx b/packages/react/src/TreeView/TreeView.features.stories.tsx
index a1e9e4a5f37..bedd9e28e48 100644
--- a/packages/react/src/TreeView/TreeView.features.stories.tsx
+++ b/packages/react/src/TreeView/TreeView.features.stories.tsx
@@ -996,14 +996,18 @@ export const LeadingAction: Story = () => {
   return (
     <TreeView aria-label="Issues">
       <TreeView.Item id="item-0">
-        <TreeView.LeadingAction icon={GrabberIcon} aria-label="Reorder item 1" />
+        <TreeView.LeadingAction>
+          <IconButton icon={GrabberIcon} aria-label="Reorder item 1" />
+        </TreeView.LeadingAction>
         <TreeView.LeadingVisual>
           <Octicon icon={IssueClosedIcon} sx={{color: 'done.fg'}} />
         </TreeView.LeadingVisual>
         Item 1
       </TreeView.Item>
       <TreeView.Item id="item-2">
-        <TreeView.LeadingAction icon={GrabberIcon} aria-label="Reorder item 2" />
+        <TreeView.LeadingAction>
+          <IconButton icon={GrabberIcon} aria-label="Reorder item 2" />
+        </TreeView.LeadingAction>
         <TreeView.LeadingVisual>
           <Octicon icon={IssueOpenedIcon} sx={{color: 'open.fg'}} />
         </TreeView.LeadingVisual>
@@ -1024,7 +1028,9 @@ export const LeadingAction: Story = () => {
         </TreeView.SubTree>
       </TreeView.Item>
       <TreeView.Item id="item-3">
-        <TreeView.LeadingAction icon={GrabberIcon} aria-label="Reorder item 3" />
+        <TreeView.LeadingAction>
+          <IconButton icon={GrabberIcon} aria-label="Reorder item 3" />
+        </TreeView.LeadingAction>
         <TreeView.LeadingVisual>
           <Octicon icon={IssueOpenedIcon} sx={{color: 'open.fg'}} />
         </TreeView.LeadingVisual>
diff --git a/packages/react/src/TreeView/TreeView.tsx b/packages/react/src/TreeView/TreeView.tsx
index b6e15f6c885..cce4e904f0f 100644
--- a/packages/react/src/TreeView/TreeView.tsx
+++ b/packages/react/src/TreeView/TreeView.tsx
@@ -21,7 +21,6 @@ import sx from '../sx'
 import {getAccessibleName} from './shared'
 import {getFirstChildElement, useRovingTabIndex} from './useRovingTabIndex'
 import {useTypeahead} from './useTypeahead'
-import {IconButton, type IconButtonProps} from '../Button'
 
 // ----------------------------------------------------------------------------
 // Context
@@ -209,7 +208,6 @@ const UlBox = styled.ul<SxProp>`
   }
 
   .PRIVATE_TreeView-item-leading-action {
-    display: flex;
     color: ${get('colors.fg.muted')};
     grid-area: leadingAction;
   }
@@ -850,16 +848,23 @@ TrailingVisual.displayName = 'TreeView.TrailingVisual'
 // ----------------------------------------------------------------------------
 // TreeView.LeadingAction
 
-const LeadingAction: React.FC<IconButtonProps> = props => {
+// ----------------------------------------------------------------------------
+// TreeView.LeadingAction
+
+const LeadingAction: React.FC<TreeViewVisualProps> = props => {
+  const {isExpanded, leadingVisualId} = React.useContext(ItemContext)
+  const children = typeof props.children === 'function' ? props.children({isExpanded}) : props.children
   return (
-    <div className="PRIVATE_TreeView-item-leading-action">
-      <IconButton variant="invisible" {...props} />
-    </div>
+    <>
+      <div className="PRIVATE_VisuallyHidden" aria-hidden={true} id={leadingVisualId}>
+        {props.label}
+      </div>
+      <div className="PRIVATE_TreeView-item-leading-action">{children}</div>
+    </>
   )
 }
 
 LeadingAction.displayName = 'TreeView.LeadingAction'
-
 // ----------------------------------------------------------------------------
 // TreeView.DirectoryIcon