Support item custom rendering

Author: Meet-student

src/Divider.tsx

@@ -6,19 +6,25 @@ import type { MenuDividerType } from './interface';
 
 export type DividerProps = Omit<MenuDividerType, 'type'>;
 
-export default function Divider({ className, style }: DividerProps) {
+export default function Divider({ className, style, itemRender }: DividerProps) {
   const { prefixCls } = React.useContext(MenuContext);
   const measure = useMeasure();
 
   if (measure) {
     return null;
   }
 
-  return (
+  const renderNode = (
     <li
       role="separator"
       className={classNames(`${prefixCls}-item-divider`, className)}
       style={style}
     />
   );
+
+  if (typeof itemRender === 'function') {
+    return itemRender(renderNode);
+  }
+
+  return renderNode;
 }

src/Menu.tsx

@@ -62,8 +62,6 @@ export interface MenuProps
   /** @deprecated Please use `items` instead */
   children?: React.ReactNode;
 
-  itemRender?: (originalNode: React.ReactNode, item?: NonNullable<ItemType>) => React.ReactElement;
-
   disabled?: boolean;
   /** @private Disable auto overflow. Pls note the prop name may refactor since we do not final decided. */
   disabledOverflow?: boolean;
@@ -160,6 +158,8 @@ export interface MenuProps
   _internalComponents?: Components;
 
   popupRender?: PopupRender;
+
+  itemRender?: (originNode: React.ReactNode, item: NonNullable<ItemType>) => React.ReactNode;
 }
 
 interface LegacyMenuProps extends MenuProps {

src/MenuItem.tsx

@@ -89,6 +89,8 @@ const InternalMenuItem = React.forwardRef((props: MenuItemProps, ref: React.Ref<
 
     onFocus,
 
+    itemRender,
+
     ...restProps
   } = props;
 
@@ -238,6 +240,10 @@ const InternalMenuItem = React.forwardRef((props: MenuItemProps, ref: React.Ref<
     </LegacyMenuItem>
   );
 
+  if (typeof itemRender === 'function') {
+    renderNode = itemRender(renderNode);
+  }
+
   if (_internalRenderMenuItem) {
     renderNode = _internalRenderMenuItem(renderNode, props, { selected });
   }

src/MenuItemGroup.tsx

@@ -52,7 +52,7 @@ const InternalMenuItemGroup = React.forwardRef<HTMLLIElement, MenuItemGroupProps
 });
 
 const MenuItemGroup = React.forwardRef<HTMLLIElement, MenuItemGroupProps>((props, ref) => {
-  const { eventKey, children } = props;
+  const { eventKey, children, itemRender } = props;
   const connectedKeyPath = useFullPath(eventKey);
   const childList: React.ReactElement[] = parseChildren(children, connectedKeyPath);
 
@@ -63,7 +63,7 @@ const MenuItemGroup = React.forwardRef<HTMLLIElement, MenuItemGroupProps>((props
 
   return (
     <InternalMenuItemGroup ref={ref} {...omit(props, ['warnKey'])}>
-      {childList}
+      {typeof itemRender === 'function' ? itemRender(childList) : childList}
     </InternalMenuItemGroup>
   );
 });

src/SubMenu/index.tsx

@@ -384,7 +384,7 @@ const InternalSubMenu = React.forwardRef<HTMLLIElement, SubMenuProps>((props, re
 });
 
 const SubMenu = React.forwardRef<HTMLLIElement, SubMenuProps>((props, ref) => {
-  const { eventKey, children } = props;
+  const { eventKey, children, itemRender } = props;
 
   const connectedKeyPath = useFullPath(eventKey);
   const childList: React.ReactElement[] = parseChildren(children, connectedKeyPath);
@@ -406,12 +406,15 @@ const SubMenu = React.forwardRef<HTMLLIElement, SubMenuProps>((props, ref) => {
   let renderNode: React.ReactNode;
 
   // ======================== Render ========================
+
+  const childListNode = typeof itemRender === 'function' ? itemRender(childList) : childList;
+
   if (measure) {
-    renderNode = childList;
+    renderNode = childListNode;
   } else {
     renderNode = (
       <InternalSubMenu ref={ref} {...props}>
-        {childList}
+        {childListNode}
       </InternalSubMenu>
     );
   }

src/interface.ts

@@ -6,6 +6,7 @@ interface ItemSharedProps {
   ref?: React.Ref<HTMLLIElement | null>;
   style?: React.CSSProperties;
   className?: string;
+  itemRender?: (originNode: React.ReactNode) => React.ReactNode;
 }
 
 export interface SubMenuType extends ItemSharedProps {

src/utils/commonUtil.ts

@@ -1,11 +1,7 @@
 import toArray from '@rc-component/util/lib/Children/toArray';
 import * as React from 'react';
 
-export function parseChildren(
-  children: React.ReactNode | undefined,
-  keyPath: string[],
-  itemRender?: (originNode: React.ReactNode) => React.ReactElement,
-) {
+export function parseChildren(children: React.ReactNode | undefined, keyPath: string[]) {
   return toArray(children).map((child, index) => {
     if (React.isValidElement(child)) {
       const { key } = child;
@@ -17,14 +13,15 @@ export function parseChildren(
         eventKey = `tmp_key-${[...keyPath, index].join('-')}`;
       }
 
-      const cloneProps = { key: eventKey, eventKey } as any;
+      const cloneProps = {
+        key: eventKey,
+        eventKey,
+      } as any;
 
       if (process.env.NODE_ENV !== 'production' && emptyKey) {
         cloneProps.warnKey = true;
       }
-      if (typeof itemRender === 'function') {
-        return itemRender(React.cloneElement(child, cloneProps));
-      }
+
       return React.cloneElement(child, cloneProps);
     }
 

src/utils/nodeUtil.tsx

@@ -21,48 +21,58 @@ function convertItemsToNodes(
 
   return (list || [])
     .map((opt, index) => {
+      const renderNodeWrapper = node => {
+        return typeof itemRender === 'function' ? itemRender(node, opt as any) : node;
+      };
       if (opt && typeof opt === 'object') {
         const { label, children, key, type, extra, ...restProps } = opt as any;
         const mergedKey = key ?? `tmp-${index}`;
 
-        let originNode: React.ReactNode = null;
-
-        // MenuItemGroup & SubMenu
+        // MenuItemGroup & SubMenuItem
         if (children || type === 'group') {
           if (type === 'group') {
-            originNode = (
-              <MergedMenuItemGroup key={mergedKey} {...restProps} title={label}>
-                {convertItemsToNodes(children, components, prefixCls, itemRender)}
+            // Group
+            return (
+              <MergedMenuItemGroup
+                key={mergedKey}
+                {...restProps}
+                itemRender={renderNodeWrapper}
+                title={label}
+              >
+                {convertItemsToNodes(children, components, prefixCls)}
               </MergedMenuItemGroup>
             );
-          } else {
-            originNode = (
-              <MergedSubMenu key={mergedKey} {...restProps} title={label}>
-                {convertItemsToNodes(children, components, prefixCls, itemRender)}
-              </MergedSubMenu>
-            );
           }
-        }
-        // Divider
-        else if (type === 'divider') {
-          originNode = <MergedDivider key={mergedKey} {...restProps} />;
-        }
-        // MenuItem
-        else {
-          originNode = (
-            <MergedMenuItem key={mergedKey} {...restProps} extra={extra}>
-              {label}
-              {(!!extra || extra === 0) && (
-                <span className={`${prefixCls}-item-extra`}>{extra}</span>
-              )}
-            </MergedMenuItem>
+
+          // Sub Menu
+          return (
+            <MergedSubMenu
+              key={mergedKey}
+              {...restProps}
+              itemRender={renderNodeWrapper}
+              title={label}
+            >
+              {convertItemsToNodes(children, components, prefixCls)}
+            </MergedSubMenu>
           );
         }
 
-        if (typeof itemRender === 'function') {
-          return itemRender(originNode, opt);
+        // MenuItem & Divider
+        if (type === 'divider') {
+          return <MergedDivider key={mergedKey} {...restProps} itemRender={renderNodeWrapper} />;
         }
-        return originNode;
+
+        return (
+          <MergedMenuItem
+            key={mergedKey}
+            {...restProps}
+            extra={extra}
+            itemRender={renderNodeWrapper}
+          >
+            {label}
+            {(!!extra || extra === 0) && <span className={`${prefixCls}-item-extra`}>{extra}</span>}
+          </MergedMenuItem>
+        );
       }
 
       return null;
@@ -92,5 +102,5 @@ export function parseItems(
     childNodes = convertItemsToNodes(items, mergedComponents, prefixCls, itemRender);
   }
 
-  return parseChildren(childNodes, keyPath, itemRender);
+  return parseChildren(childNodes, keyPath);
 }