allow chaining polymorphic components

packages/itwinui-react/src/core/Menu/MenuItem.tsx

@@ -8,12 +8,13 @@ import {
   useMergedRefs,
   useId,
   createWarningLogger,
+  ButtonBase,
+  polymorphic,
 } from '../../utils/index.js';
 import type { PolymorphicForwardRefComponent } from '../../utils/index.js';
 import { Menu, MenuContext } from './Menu.js';
 import { ListItem } from '../List/ListItem.js';
 import type { ListItemOwnProps } from '../List/ListItem.js';
-import cx from 'classnames';
 
 const logWarning = createWarningLogger();
 
@@ -83,7 +84,6 @@ export type MenuItemProps = {
  */
 export const MenuItem = React.forwardRef((props, forwardedRef) => {
   const {
-    className,
     children,
     isSelected,
     disabled,
@@ -129,26 +129,15 @@ export const MenuItem = React.forwardRef((props, forwardedRef) => {
     } satisfies Parameters<typeof Menu>[0]['popoverProps'];
   }, [subMenuItems.length]);
 
-  const onClick = () => {
-    if (disabled) {
-      return;
-    }
-
-    onClickProp?.(value);
-  };
-
-  const handlers: React.DOMAttributes<HTMLButtonElement> = { onClick };
+  const onClick = () => onClickProp?.(value);
 
   /** Index of this item out of all the focusable items in the parent `Menu` */
   const focusableItemIndex = parentMenu?.focusableElements.findIndex(
     (el) => el === menuItemRef.current,
   );
 
   const trigger = (
-    <ListItem
-      as='button'
-      type='button'
-      className={cx('iui-button-base', className)}
+    <MenuItemBase
       actionable
       size={size}
       active={isSelected}
@@ -163,10 +152,10 @@ export const MenuItem = React.forwardRef((props, forwardedRef) => {
       {...(parentMenu?.popoverGetItemProps != null
         ? parentMenu.popoverGetItemProps({
             focusableItemIndex,
-            userProps: handlers,
+            userProps: { onClick },
           })
-        : handlers)}
-      {...(rest as React.DOMAttributes<HTMLButtonElement>)}
+        : { onClick })}
+      {...(rest as React.DOMAttributes<HTMLElement>)}
     >
       {startIcon && (
         <ListItem.Icon as='span' aria-hidden>
@@ -187,7 +176,7 @@ export const MenuItem = React.forwardRef((props, forwardedRef) => {
           {endIcon}
         </ListItem.Icon>
       )}
-    </ListItem>
+    </MenuItemBase>
   );
 
   return (
@@ -206,6 +195,11 @@ if (process.env.NODE_ENV === 'development') {
   MenuItem.displayName = 'MenuItem';
 }
 
+const MenuItemBase = polymorphic(
+  ButtonBase,
+  ListItem,
+) as PolymorphicForwardRefComponent<'button', ListItemOwnProps>;
+
 // ----------------------------------------------------------------------------
 
 export type TreeEvent = {

packages/itwinui-react/src/utils/functions/polymorphic.test.tsx

@@ -4,6 +4,32 @@
  *--------------------------------------------------------------------------------------------*/
 import { render } from '@testing-library/react';
 import { polymorphic } from './polymorphic.js';
+import type { PolymorphicForwardRefComponent } from '../props.js';
+
+it('should work when called directly', () => {
+  const MyDiv1 = polymorphic.div('my-div1');
+  const MyDiv2 = polymorphic.div('my-div2');
+  const MyDiv = polymorphic(
+    MyDiv1,
+    MyDiv2,
+  ) as PolymorphicForwardRefComponent<'div'>;
+
+  const { container: container1 } = render(<MyDiv data-testid='foo'>🍏</MyDiv>);
+  const el1 = container1.querySelector('div') as HTMLElement;
+  expect(el1).toHaveClass('my-div1', 'my-div2');
+  expect(el1).toHaveTextContent('🍏');
+  expect(el1).toHaveAttribute('data-testid', 'foo');
+
+  const { container: container2 } = render(
+    <MyDiv as='span' data-testid='bar'>
+      🥭
+    </MyDiv>,
+  );
+  const el2 = container2.querySelector('span') as HTMLElement;
+  expect(el2).toHaveClass('my-div1', 'my-div2');
+  expect(el2).toHaveTextContent('🥭');
+  expect(el2).toHaveAttribute('data-testid', 'bar');
+});
 
 it('should work when called as property', () => {
   const MyButton = polymorphic.button('my-button');

packages/itwinui-react/src/utils/functions/polymorphic.tsx

@@ -8,7 +8,7 @@ import type { PolymorphicForwardRefComponent } from '../props.js';
 import { useGlobals } from '../hooks/useGlobals.js';
 import { styles } from '../../styles.js';
 
-const _base = <As extends keyof JSX.IntrinsicElements = 'div'>(
+const _elementBase = <As extends keyof JSX.IntrinsicElements = 'div'>(
   defaultElement: As,
 ) => {
   return (className: string, attrs?: JSX.IntrinsicElements[As]) => {
@@ -21,6 +21,16 @@ const _base = <As extends keyof JSX.IntrinsicElements = 'div'>(
         ),
       };
 
+      useGlobals();
+
+      if (Array.isArray(as as any)) {
+        if (as.length > 1) {
+          const [Component, ...restAs] = as;
+          return <Component ref={ref} {...props} as={restAs} />;
+        }
+        (as as any) = as[0];
+      }
+
       const Element = (as as any) || 'div';
 
       // Add tabIndex to interactive elements if not already set.
@@ -33,8 +43,6 @@ const _base = <As extends keyof JSX.IntrinsicElements = 'div'>(
         props.tabIndex ??= 0;
       }
 
-      useGlobals();
-
       return <Element ref={ref} {...props} />;
     }) as PolymorphicForwardRefComponent<NonNullable<typeof defaultElement>>;
 
@@ -46,6 +54,19 @@ const _base = <As extends keyof JSX.IntrinsicElements = 'div'>(
   };
 };
 
+const _componentBase = (
+  ...components: PolymorphicForwardRefComponent<any>[]
+) => {
+  const Comp = _elementBase('div')('');
+  return React.forwardRef((props, ref) => (
+    <Comp
+      ref={ref}
+      {...props}
+      as={[...components, props.as].filter(Boolean) as any}
+    />
+  )) as PolymorphicForwardRefComponent<'div'>;
+};
+
 /**
  * Utility to create a type-safe polymorphic component with a simple class.
  *
@@ -57,23 +78,31 @@ const _base = <As extends keyof JSX.IntrinsicElements = 'div'>(
  * - adds and merges CSS classes
  * - adds tabIndex to interactive elements (Safari workaround)
  *
+ * Alternatively, it can be called directly with multiple components to create a chain of
+ * components. This is useful for creating wrapper components that still need to support
+ * user-specified `as` prop. **Note:** All components in the chain must be polymorphic.
+ *
  * @example
  * const MyPolyButton = polymorphic.button('my-poly-button', { type: 'button' });
  * <MyPolyButton as='a' href='#'>...</MyPolyButton>;
  *
+ * @example
+ * const MyMenuButton = polymorphic(MyPolyButton, MyMenu) as PolymorphicForwardRefComponent<'button'>;
+ * <MyMenuButton as='a' href='#'>...</MyMenuButton>;
+ *
  * @private
  */
-export const polymorphic = new Proxy({} as never, {
+export const polymorphic = new Proxy(_componentBase, {
   get: (target, prop) => {
     if (typeof prop === 'string') {
       // eslint-disable-next-line -- string is as far as we can narrow it down
       // @ts-ignore
-      return _base(prop);
+      return _elementBase(prop);
     }
     return Reflect.get(target, prop);
   },
-}) as {
-  [key in keyof JSX.IntrinsicElements]: ReturnType<typeof _base<key>>;
+}) as typeof _componentBase & {
+  [key in keyof JSX.IntrinsicElements]: ReturnType<typeof _elementBase<key>>;
 };
 
 // e.g. iui-list-item-icon -> ListItemIcon