chore: forward ref

Author: segunadebayo

packages/react/src/components/json-tree-view/json-tree-view-root-provider.tsx

@@ -1,8 +1,11 @@
 import type { JsonNode } from '@zag-js/json-tree-utils'
+import { forwardRef } from 'react'
 import { TreeView } from '../tree-view'
 
 export interface JsonTreeViewRootProviderProps extends TreeView.RootProviderProps<JsonNode> {}
 
-export const JsonTreeViewRootProvider = (props: JsonTreeViewRootProviderProps) => {
-  return <TreeView.RootProvider data-scope="json-tree-view" {...props} />
-}
+export const JsonTreeViewRootProvider = forwardRef<HTMLDivElement, JsonTreeViewRootProviderProps>((props, ref) => {
+  return <TreeView.RootProvider data-scope="json-tree-view" {...props} ref={ref} />
+})
+
+JsonTreeViewRootProvider.displayName = 'JsonTreeViewRootProvider'

packages/react/src/components/json-tree-view/json-tree-view-root.tsx

@@ -1,12 +1,12 @@
 import { type JsonNode, getRootNode, nodeToString, nodeToValue } from '@zag-js/json-tree-utils'
-import { useMemo } from 'react'
+import { forwardRef, useMemo } from 'react'
 import { TreeView, createTreeCollection } from '../tree-view'
 
-export interface JsonTreeViewRootProps<T = unknown> extends Omit<TreeView.RootProps<JsonNode>, 'collection'> {
-  data: T
+export interface JsonTreeViewRootProps extends Omit<TreeView.RootProps<JsonNode>, 'collection'> {
+  data: unknown
 }
 
-export function JsonTreeViewRoot<T = unknown>(props: JsonTreeViewRootProps<T>) {
+export const JsonTreeViewRoot = forwardRef<HTMLDivElement, JsonTreeViewRootProps>((props, ref) => {
   const { data, ...restProps } = props
 
   const collection = useMemo(() => {
@@ -17,5 +17,9 @@ export function JsonTreeViewRoot<T = unknown>(props: JsonTreeViewRootProps<T>) {
     })
   }, [data])
 
-  return <TreeView.Root data-scope="json-tree-view" typeahead={false} collection={collection} {...restProps} />
-}
+  return (
+    <TreeView.Root data-scope="json-tree-view" collection={collection} {...restProps} ref={ref} typeahead={false} />
+  )
+})
+
+JsonTreeViewRoot.displayName = 'JsonTreeViewRoot'

packages/react/src/components/json-tree-view/json-tree-view-tree.tsx

@@ -1,18 +1,21 @@
+import { forwardRef } from 'react'
 import { createSplitProps } from '../../utils/create-split-props'
 import { TreeView, useTreeViewContext } from '../tree-view'
 import { JsonTreeViewNode, type JsonTreeViewNodeBaseProps } from './json-tree-view-node'
 
 export interface JsonTreeViewTreeProps extends TreeView.TreeProps, JsonTreeViewNodeBaseProps {}
 
-export const JsonTreeViewTree = (props: JsonTreeViewTreeProps) => {
+export const JsonTreeViewTree = forwardRef<HTMLDivElement, JsonTreeViewTreeProps>((props, ref) => {
   const [nodeProps, treeProps] = createSplitProps<JsonTreeViewNodeBaseProps>()(props, ['arrowIcon', 'showIndentGuide'])
   const tree = useTreeViewContext()
   const children = tree.collection.getNodeChildren(tree.collection.rootNode)
   return (
-    <TreeView.Tree data-scope="json-tree-view" {...treeProps}>
+    <TreeView.Tree data-scope="json-tree-view" {...treeProps} ref={ref}>
       {children.map((child, index) => (
         <JsonTreeViewNode key={child.id} node={child} indexPath={[index]} {...nodeProps} />
       ))}
     </TreeView.Tree>
   )
-}
+})
+
+JsonTreeViewTree.displayName = 'JsonTreeViewTree'

packages/react/src/components/tree-view/use-tree-view.ts

@@ -6,7 +6,7 @@ import type { Optional } from '../../types'
 import type { TreeCollection, TreeNode } from '../collection'
 
 export interface UseTreeViewProps<T extends TreeNode>
-  extends Optional<Omit<treeView.Props, 'dir' | 'getRootNode' | 'collection'>, 'id'> {
+  extends Optional<Omit<treeView.Props<T>, 'dir' | 'getRootNode' | 'collection'>, 'id'> {
   /**
    * The collection of tree nodes
    */