Feat: NewRefProp

packages/@headlessui-react/src/components/menu/menu.tsx

@@ -70,6 +70,7 @@ import {
   render,
   useMergeRefsFn,
   type HasDisplayName,
+  type NewRefProp,
   type RefProp,
 } from '../../utils/render'
 import { useDescriptions } from '../description/description'
@@ -1096,7 +1097,7 @@ function SeparatorFn<TTag extends ElementType = typeof DEFAULT_SEPARATOR_TAG>(
 
 export interface _internal_ComponentMenu extends HasDisplayName {
   <TTag extends ElementType = typeof DEFAULT_MENU_TAG>(
-    props: MenuProps<TTag> & RefProp<typeof MenuFn>
+    props: MenuProps<TTag> & NewRefProp<TTag>
   ): JSX.Element
 }
 

packages/@headlessui-react/src/utils/render.ts

@@ -6,6 +6,7 @@ import {
   isValidElement,
   useCallback,
   useRef,
+  type ComponentPropsWithRef,
   type ElementType,
   type MutableRefObject,
   type ReactElement,
@@ -386,6 +387,10 @@ export type RefProp<T extends Function> = T extends (props: any, ref: Ref<infer
   ? { ref?: Ref<RefType> }
   : never
 
+export type NewRefProp<T extends ElementType> = unknown extends ComponentPropsWithRef<T>['ref']
+  ? {}
+  : { ref?: ComponentPropsWithRef<T>['ref'] }
+
 // TODO: add proper return type, but this is not exposed as public API so it's fine for now
 export function mergeProps<T extends Props<any, any>[]>(...listOfProps: T) {
   if (listOfProps.length === 0) return {}

playgrounds/react/pages/menu/test-menu.tsx

@@ -0,0 +1,60 @@
+// TODO: delete this file
+
+import { ComponentPropsWithRef, ElementType, Fragment, forwardRef, useRef } from 'react'
+
+// Example components that `Menu` will be used `as`
+
+const MenuWithRef = ({}: { ref: React.RefObject<HTMLButtonElement> }) => null
+const MenuWithoutRef = ({}: {}) => null
+const MenuWithCustomRef = ({}: { ref: React.RefObject<{ onOpen?: () => void }> }) => null
+const MenuWithForwardRef = forwardRef<HTMLButtonElement, {}>((props, ref) => null)
+const MenuWithForwardRefWithCustomRef = forwardRef<{ onOpen: () => void }, {}>((props, ref) => null)
+
+/**
+ * The `as` prop can be a few things:
+ * - A regular HTML tag
+ * - Another component
+ * - A React fragment
+ *
+ * So, the ref can also point to different things:
+ * - A reference to an HTML tag
+ * - A ForwardRef, which could be for an HTML tag or a custom one
+ * - A component with a `ref` prop (React v19+), pointing to either an HTML tag or a custom ref
+ */
+
+type RefProp<T extends ElementType> = unknown extends ComponentPropsWithRef<T>['ref']
+  ? {}
+  : { ref?: ComponentPropsWithRef<T>['ref'] }
+
+type Menu = {
+  <TTag extends ElementType = 'div'>(
+    props: {
+      as?: TTag
+    } & RefProp<TTag>
+  ): JSX.Element
+}
+
+const MenuAs = forwardRef(({ as }: any, ref) => {
+  const Component = as || 'div'
+
+  return <Component ref={ref} />
+}) as Menu
+
+export default function MenuExample() {
+  const divRef = useRef<HTMLDivElement>(null)
+  const buttonRef = useRef<HTMLButtonElement>(null)
+  const customRef = useRef<{ onOpen: () => void }>(null)
+
+  return (
+    <Fragment>
+      <MenuAs ref={divRef} /> {/* Default `as` is a div */}
+      <MenuAs as="button" ref={buttonRef} />
+      <MenuAs as={Fragment} /> {/* ref is `never` */}
+      <MenuAs as={MenuWithRef} ref={buttonRef} />
+      <MenuAs as={MenuWithoutRef} /> {/* ref is `never` */}
+      <MenuAs as={MenuWithCustomRef} ref={customRef} />
+      <MenuAs as={MenuWithForwardRef} ref={buttonRef} />
+      <MenuAs as={MenuWithForwardRefWithCustomRef} ref={customRef} />
+    </Fragment>
+  )
+}