From 484986dae84204a2e092e01ffff46f2f352b11a7 Mon Sep 17 00:00:00 2001
From: MadCcc <1075746765@qq.com>
Date: Sun, 27 Mar 2022 15:32:22 +0800
Subject: [PATCH] feat: menu ref support focus

---
 src/Menu.tsx       | 15 +++++++++++----
 src/interface.ts   |  5 +++++
 tests/Menu.spec.js | 17 +++++++++++++++--
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/Menu.tsx b/src/Menu.tsx
index d671124c..49d1e0b2 100644
--- a/src/Menu.tsx
+++ b/src/Menu.tsx
@@ -15,6 +15,7 @@ import type {
   SelectInfo,
   RenderIconType,
   ItemType,
+  MenuRef,
 } from './interface';
 import MenuItem from './MenuItem';
 import { parseItems } from './utils/nodeUtil';
@@ -28,7 +29,7 @@ import { PathRegisterContext, PathUserContext } from './context/PathContext';
 import useKeyRecords, { OVERFLOW_KEY } from './hooks/useKeyRecords';
 import { IdContext } from './context/IdContext';
 import PrivateContext from './context/PrivateContext';
-import { composeRef } from 'rc-util/lib/ref';
+import { useImperativeHandle } from 'react';
 
 /**
  * Menu modify after refactor:
@@ -149,7 +150,7 @@ interface LegacyMenuProps extends MenuProps {
   openAnimation?: string;
 }
 
-const Menu = React.forwardRef<HTMLUListElement, MenuProps>((props, ref) => {
+const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
   const {
     prefixCls = 'rc-menu',
     style,
@@ -232,7 +233,13 @@ const Menu = React.forwardRef<HTMLUListElement, MenuProps>((props, ref) => {
   const [mounted, setMounted] = React.useState(false);
 
   const containerRef = React.useRef<HTMLUListElement>();
-  const mergedRef = composeRef(containerRef, ref);
+
+  useImperativeHandle(ref, () => ({
+    list: containerRef.current,
+    focus: (options?: FocusOptions) => {
+      containerRef.current?.focus(options);
+    },
+  }));
 
   const uuid = useUUID(id);
 
@@ -498,7 +505,7 @@ const Menu = React.forwardRef<HTMLUListElement, MenuProps>((props, ref) => {
   const container = (
     <Overflow
       id={id}
-      ref={mergedRef as any}
+      ref={containerRef as any}
       prefixCls={`${prefixCls}-overflow`}
       component="ul"
       itemComponent={MenuItem}
diff --git a/src/interface.ts b/src/interface.ts
index 04d18bd2..92e2ac7c 100644
--- a/src/interface.ts
+++ b/src/interface.ts
@@ -116,3 +116,8 @@ export type SelectEventHandler = (info: SelectInfo) => void;
 
 // ========================== Click ==========================
 export type MenuClickEventHandler = (info: MenuInfo) => void;
+
+export type MenuRef = {
+  focus: (options?: FocusOptions) => void;
+  list: HTMLUListElement;
+};
diff --git a/tests/Menu.spec.js b/tests/Menu.spec.js
index 430fa1b7..d5e2d047 100644
--- a/tests/Menu.spec.js
+++ b/tests/Menu.spec.js
@@ -650,14 +650,27 @@ describe('Menu', () => {
 
   it('should support ref', () => {
     const menuRef = React.createRef();
+    const wrapper = mount(
+      <Menu ref={menuRef}>
+        <MenuItem key="light">Light</MenuItem>
+      </Menu>,
+    );
 
-    mount(
+    expect(menuRef.current?.list).toBe(wrapper.find('ul').first().getDOMNode());
+  });
+
+  it('should support focus through ref', () => {
+    const menuRef = React.createRef();
+    const wrapper = mount(
       <Menu ref={menuRef}>
         <MenuItem key="light">Light</MenuItem>
       </Menu>,
     );
+    menuRef.current?.focus();
 
-    expect(menuRef.current).toBeTruthy();
+    expect(document.activeElement).toBe(
+      wrapper.find('ul').first().getDOMNode(),
+    );
   });
 });
 /* eslint-enable */