Example.tsx

import { ChevronUpDownIcon } from '@heroicons/react/24/outline';
import { PropsWithChildren, useContext, useState } from 'react';
import {
  Button as ButtonBase,
  Collection,
  Popover, SelectStateContext,
  TreeItemProps,
  UNSTABLE_Tree as Tree,
  UNSTABLE_TreeItem as TreeItem,
  UNSTABLE_TreeItemContent as TreeItemContent,
} from 'react-aria-components';
import { useAsyncList } from 'react-stately';
import { ChevronRightIcon } from '@heroicons/react/20/solid';
import clsx from 'clsx';
import { SelectedValueContext, TreeSelect, SelectValue } from '@/components/TreeSelect';
import { Button, Label } from '@/components/common';
import { loadItemChildren } from './documents';

function TreeNode({
  id, children, hasChildren, ...props
}: PropsWithChildren<TreeItemProps & {
  hasChildren: boolean
}>) {
  const selectState = useContext(SelectStateContext);
  const selectedValue = useContext(SelectedValueContext);
  const [isExpanded, setExpanded] = useState(false);

  const itemsList = useAsyncList({
    load: async () => {
      const items = await loadItemChildren(id as string);
      return { items };
    },
  });

  return (
    <TreeItem childItems={isExpanded ? itemsList.items : []} {...props}>
      <TreeItemContent>
        {({ level }) => (
          <ButtonBase
            className="text-sm w-full group flex items-center gap-2 cursor-default select-none py-2 px-1 outline-none text-gray-900 focus:bg-gray-100 hover:bg-gray-100 data-[hovered]:bg-gray-100"
            {...hasChildren ? {
              onPress: () => {
                setExpanded((s) => !s);
              },
              slot: 'chevron',
            } : {
              onPress: () => {
                selectedValue?.setValue(props.textValue);
                selectState.setSelectedKey(id as any);
                selectState.close();
              },
            }}
          >
            <div style={{ width: `${(level - 1) * 15}px` }} />
            <div className="pl-5 flex items-center relative">
              {hasChildren && (
                <ChevronRightIcon
                  className={clsx('h-4 w-4 absolute left-0 transition-transform duration-150', isExpanded && 'rotate-90')}
                />
              )}
              {children}
            </div>
          </ButtonBase>
        )}
      </TreeItemContent>
      <Collection items={itemsList.items}>
        {(item: any) => (
          <TreeNode
            id={item.id}
            textValue={item.name}
            hasChildren={item.hasChildren}
          >
            {item.name}
          </TreeNode>
        )}
      </Collection>
    </TreeItem>
  );
}

export default function FilePicker() {
  const rootList = useAsyncList({
    load: async () => {
      const items = await loadItemChildren('1');
      return { items };
    },
  });

  return (
    <TreeSelect onSelectionChange={(k) => console.log(k)}>
      <Label slot="label">Document</Label>
      <Button
        className="w-full flex border border-gray-300 hover:border-gray-400 items-center cursor-pointer
                         rounded-md bg-white transition py-2 pl-4 pr-2 text-left
                         text-gray-700 focus:outline-none focus-visible:ring-2 ring-brand-400/30 ring-offset-1
                         disabled:bg-gray-50 disabled:text-gray-400 disabled:ring-gray-200 disabled:!opacity-100"
      >
        <SelectValue className="flex-1 truncate data-[placeholder]:text-gray-500 text-sm font-normal" />
        <ChevronUpDownIcon
          className="h-5 w-5 text-gray-400"
          aria-hidden="true"
        />
      </Button>
      <Popover className="react-aria-Popover w-[--trigger-width]">
        <Tree
          className="react-aria-Tree -mx-4 -my-2 text-gray-800 focus:bg-transparent outline-none"
          aria-label="Tree"
          items={rootList.items}
        >
          {(item) => (
            <TreeNode hasChildren={item.hasChildren} id={item.id} textValue={item.name}>
              {item.name}
            </TreeNode>
          )}
        </Tree>
      </Popover>
    </TreeSelect>
  );
}

TreeSelect.tsx

import { useResizeObserver } from '@react-aria/utils';
import { useFocusRing, useSelect } from 'react-aria';
import {
  createContext,
  Dispatch,
  ForwardedRef,
  forwardRef,
  ReactNode,
  SetStateAction,
  useCallback,
  useContext,
  useRef,
  useState,
} from 'react';
import {
  ButtonContext,
  LabelContext,
  OverlayTriggerStateContext,
  PopoverContext,
  Provider,
  SelectContext,
  SelectStateContext,
  SelectValueContext,
  SelectValueProps,
  TextContext,
  useContextProps,
  useSlottedContext,
} from 'react-aria-components';
import {SelectStateOptions, useOverlayTriggerState, useSelectState} from 'react-stately';
import { useSlot } from './slot';
import { useRenderProps } from './common/render-props';
import { forwardRefType } from './common/types';

export const SelectedValueContext = createContext<{
  value: string | null,
  setValue: Dispatch<SetStateAction<string | null>>
} | null>(null);

export function TreeSelect<T extends object>({ children, ...selectProps }: { children?: ReactNode | undefined } & SelectStateOptions<T>) {
  [selectProps] = useContextProps(selectProps, null, SelectContext);

  const state = useSelectState(selectProps);
  const overlayState = useOverlayTriggerState({});
  const [selectedValue, setSelectedValue] = useState<string | null>(null);

  (state as any).isOpen = overlayState.isOpen;
  (state as any).toggle = overlayState.toggle;
  (state as any).open = overlayState.open;
  (state as any).setOpen = overlayState.setOpen;
  (state as any).close = overlayState.close;

  const { isFocusVisible, focusProps } = useFocusRing({ within: true });

  const buttonRef = useRef<HTMLButtonElement>(null);
  const [labelRef, label] = useSlot();

  const {
    labelProps,
    triggerProps,
    valueProps,
  } = useSelect({
    label,
  }, state, buttonRef);

  const [buttonWidth, setButtonWidth] = useState<string | null>(null);
  const onResize = useCallback(() => {
    if (buttonRef.current) {
      setButtonWidth(`${buttonRef.current.offsetWidth}px`);
    }
  }, [buttonRef]);

  useResizeObserver({
    ref: buttonRef,
    onResize,
  });

  return (
    <Provider values={[
      [LabelContext, { ...labelProps, ref: labelRef, elementType: 'span' }],
      [OverlayTriggerStateContext, overlayState],
      [SelectContext, {}],
      [SelectStateContext, state],
      [SelectValueContext, valueProps],
      [ButtonContext, {
        ...triggerProps,
        ref: buttonRef,
        isPressed: state.isOpen,
      }],
      [PopoverContext, {
        trigger: 'TreeSelect',
        triggerRef: buttonRef,
        placement: 'bottom start',
        style: { '--trigger-width': buttonWidth } as React.CSSProperties,
      }],
      [SelectedValueContext, { value: selectedValue, setValue: setSelectedValue } as any],
      [TextContext, {
        slots: {
          // description: descriptionProps,
          // errorMessage: errorMessageProps,
        },
      }],
    ]}
    >
      <div
        {...focusProps}
        data-focused={state.isFocused || undefined}
        data-focus-visible={isFocusVisible || undefined}
        data-open={state.isOpen || undefined}
      >
        {children}
      </div>
    </Provider>
  );
}

function SelectValue<T extends object>(
  props: SelectValueProps<T>,
  ref: ForwardedRef<HTMLSpanElement>,
) {
  [props, ref] = useContextProps(props, ref, SelectValueContext);
  const state = useContext(SelectStateContext)!;
  const { placeholder } = useSlottedContext(SelectContext)!;
  const { value } = useContext(SelectedValueContext)!;
  const selectedItem = state.selectedKey != null
    ? value
    : null;

  const renderProps = useRenderProps({
    ...props,
    defaultChildren: selectedItem || placeholder || 'Select an Item',
    defaultClassName: 'react-aria-SelectValue',
    values: {
      selectedItem: state.selectedItem?.value as T ?? null,
      selectedText: state.selectedItem?.textValue ?? null,
      isPlaceholder: !selectedItem,
    } as any,
  } as any);

  return (
    <span ref={ref} {...props} {...renderProps} data-placeholder={!selectedItem || undefined}>
      <TextContext.Provider value={undefined}>
        {renderProps.children}
      </TextContext.Provider>
    </span>
  );
}

// eslint-disable-next-line no-underscore-dangle
const _SelectValue = (forwardRef as forwardRefType)(SelectValue);
export { _SelectValue as SelectValue };