From 3c74ad1a91c758904fa02742a5d736ab52a63f79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=8C=E8=B4=A7=E6=9C=BA=E5=99=A8=E4=BA=BA?=
 <smith3816@gmail.com>
Date: Thu, 23 Dec 2021 14:54:58 +0800
Subject: [PATCH] chore: support internal props (#428)

* chore: support internal props

* chore: more props
---
 src/Menu.tsx                  | 114 ++++++++++++++++++++--------------
 src/MenuItem.tsx              |  12 +++-
 src/context/PrivateContext.ts |  10 +++
 tests/Private.spec.js         |  26 ++++++++
 4 files changed, 116 insertions(+), 46 deletions(-)
 create mode 100644 src/context/PrivateContext.ts
 create mode 100644 tests/Private.spec.js

diff --git a/src/Menu.tsx b/src/Menu.tsx
index 376b0053..fcf59025 100644
--- a/src/Menu.tsx
+++ b/src/Menu.tsx
@@ -26,6 +26,7 @@ import useUUID from './hooks/useUUID';
 import { PathRegisterContext, PathUserContext } from './context/PathContext';
 import useKeyRecords, { OVERFLOW_KEY } from './hooks/useKeyRecords';
 import { IdContext } from './context/IdContext';
+import PrivateContext from './context/PrivateContext';
 
 /**
  * Menu modify after refactor:
@@ -110,6 +111,16 @@ export interface MenuProps
   // >>>>> Events
   onClick?: MenuClickEventHandler;
   onOpenChange?: (openKeys: string[]) => void;
+
+  // >>>>> Internal
+  /***
+   * @private Only used for `pro-layout`. Do not use in your prod directly
+   * and we do not promise any compatibility for this.
+   */
+  _internalRenderMenuItem?: (
+    originNode: React.ReactElement,
+    menuItemProps: any,
+  ) => React.ReactElement;
 }
 
 interface LegacyMenuProps extends MenuProps {
@@ -184,6 +195,9 @@ const Menu: React.FC<MenuProps> = props => {
     openAnimation,
     openTransitionName,
 
+    // Internal
+    _internalRenderMenuItem,
+
     ...restProps
   } = props as LegacyMenuProps;
 
@@ -426,6 +440,14 @@ const Menu: React.FC<MenuProps> = props => {
     setMounted(true);
   }, []);
 
+  // ======================= Context ========================
+  const privateContext = React.useMemo(
+    () => ({
+      _internalRenderMenuItem,
+    }),
+    [_internalRenderMenuItem],
+  );
+
   // ======================== Render ========================
 
   // >>>>> Children
@@ -502,51 +524,53 @@ const Menu: React.FC<MenuProps> = props => {
 
   // >>>>> Render
   return (
-    <IdContext.Provider value={uuid}>
-      <MenuContextProvider
-        prefixCls={prefixCls}
-        mode={mergedMode}
-        openKeys={mergedOpenKeys}
-        rtl={isRtl}
-        // Disabled
-        disabled={disabled}
-        // Motion
-        motion={mounted ? motion : null}
-        defaultMotions={mounted ? defaultMotions : null}
-        // Active
-        activeKey={mergedActiveKey}
-        onActive={onActive}
-        onInactive={onInactive}
-        // Selection
-        selectedKeys={mergedSelectKeys}
-        // Level
-        inlineIndent={inlineIndent}
-        // Popup
-        subMenuOpenDelay={subMenuOpenDelay}
-        subMenuCloseDelay={subMenuCloseDelay}
-        forceSubMenuRender={forceSubMenuRender}
-        builtinPlacements={builtinPlacements}
-        triggerSubMenuAction={triggerSubMenuAction}
-        getPopupContainer={getInternalPopupContainer}
-        // Icon
-        itemIcon={itemIcon}
-        expandIcon={expandIcon}
-        // Events
-        onItemClick={onInternalClick}
-        onOpenChange={onInternalOpenChange}
-      >
-        <PathUserContext.Provider value={pathUserContext}>
-          {container}
-        </PathUserContext.Provider>
-
-        {/* Measure menu keys. Add `display: none` to avoid some developer miss use the Menu */}
-        <div style={{ display: 'none' }} aria-hidden>
-          <PathRegisterContext.Provider value={registerPathContext}>
-            {childList}
-          </PathRegisterContext.Provider>
-        </div>
-      </MenuContextProvider>
-    </IdContext.Provider>
+    <PrivateContext.Provider value={privateContext}>
+      <IdContext.Provider value={uuid}>
+        <MenuContextProvider
+          prefixCls={prefixCls}
+          mode={mergedMode}
+          openKeys={mergedOpenKeys}
+          rtl={isRtl}
+          // Disabled
+          disabled={disabled}
+          // Motion
+          motion={mounted ? motion : null}
+          defaultMotions={mounted ? defaultMotions : null}
+          // Active
+          activeKey={mergedActiveKey}
+          onActive={onActive}
+          onInactive={onInactive}
+          // Selection
+          selectedKeys={mergedSelectKeys}
+          // Level
+          inlineIndent={inlineIndent}
+          // Popup
+          subMenuOpenDelay={subMenuOpenDelay}
+          subMenuCloseDelay={subMenuCloseDelay}
+          forceSubMenuRender={forceSubMenuRender}
+          builtinPlacements={builtinPlacements}
+          triggerSubMenuAction={triggerSubMenuAction}
+          getPopupContainer={getInternalPopupContainer}
+          // Icon
+          itemIcon={itemIcon}
+          expandIcon={expandIcon}
+          // Events
+          onItemClick={onInternalClick}
+          onOpenChange={onInternalOpenChange}
+        >
+          <PathUserContext.Provider value={pathUserContext}>
+            {container}
+          </PathUserContext.Provider>
+
+          {/* Measure menu keys. Add `display: none` to avoid some developer miss use the Menu */}
+          <div style={{ display: 'none' }} aria-hidden>
+            <PathRegisterContext.Provider value={registerPathContext}>
+              {childList}
+            </PathRegisterContext.Provider>
+          </div>
+        </MenuContextProvider>
+      </IdContext.Provider>
+    </PrivateContext.Provider>
   );
 };
 
diff --git a/src/MenuItem.tsx b/src/MenuItem.tsx
index b48ec56a..2e142924 100644
--- a/src/MenuItem.tsx
+++ b/src/MenuItem.tsx
@@ -17,6 +17,7 @@ import Icon from './Icon';
 import useDirectionStyle from './hooks/useDirectionStyle';
 import { useFullPath, useMeasure } from './context/PathContext';
 import { useMenuId } from './context/IdContext';
+import PrivateContext from './context/PrivateContext';
 
 export interface MenuItemProps
   extends Omit<
@@ -116,6 +117,9 @@ const InternalMenuItem = (props: MenuItemProps) => {
     // Active
     onActive,
   } = React.useContext(MenuContext);
+
+  const { _internalRenderMenuItem } = React.useContext(PrivateContext);
+
   const itemCls = `${prefixCls}-item`;
 
   const legacyMenuItemRef = React.useRef<any>();
@@ -199,7 +203,7 @@ const InternalMenuItem = (props: MenuItemProps) => {
     optionRoleProps['aria-selected'] = selected;
   }
 
-  return (
+  let renderNode = (
     <LegacyMenuItem
       ref={legacyMenuItemRef}
       elementRef={elementRef}
@@ -238,6 +242,12 @@ const InternalMenuItem = (props: MenuItemProps) => {
       />
     </LegacyMenuItem>
   );
+
+  if (_internalRenderMenuItem) {
+    renderNode = _internalRenderMenuItem(renderNode, props);
+  }
+
+  return renderNode;
 };
 
 function MenuItem(props: MenuItemProps): React.ReactElement {
diff --git a/src/context/PrivateContext.ts b/src/context/PrivateContext.ts
new file mode 100644
index 00000000..2b580c6f
--- /dev/null
+++ b/src/context/PrivateContext.ts
@@ -0,0 +1,10 @@
+import * as React from 'react';
+import type { MenuProps } from '../Menu';
+
+export interface PrivateContextProps {
+  _internalRenderMenuItem?: MenuProps['_internalRenderMenuItem'];
+}
+
+const PrivateContext = React.createContext<PrivateContextProps>({});
+
+export default PrivateContext;
diff --git a/tests/Private.spec.js b/tests/Private.spec.js
new file mode 100644
index 00000000..2ba723e9
--- /dev/null
+++ b/tests/Private.spec.js
@@ -0,0 +1,26 @@
+/* eslint-disable no-undef */
+import React from 'react';
+import { mount } from 'enzyme';
+import classnames from 'classnames';
+import Menu, { MenuItem } from '../src';
+
+describe('Private Props', () => {
+  it('_internalRenderMenuItem', () => {
+    const wrapper = mount(
+      <Menu
+        _internalRenderMenuItem={node =>
+          React.cloneElement(node, {
+            className: classnames(node.props.className, 'inject-cls'),
+          })
+        }
+      >
+        <MenuItem key="1">1</MenuItem>
+      </Menu>,
+    );
+
+    console.log(wrapper.html());
+
+    expect(wrapper.exists('.inject-cls')).toBeTruthy();
+  });
+});
+/* eslint-enable */