Fixed group and submenu focus

src/Menu.tsx

@@ -2,16 +2,20 @@ import classNames from 'classnames';
 import type { CSSMotionProps } from 'rc-motion';
 import Overflow from 'rc-overflow';
 import useMergedState from 'rc-util/lib/hooks/useMergedState';
+import isEqual from 'rc-util/lib/isEqual';
 import warning from 'rc-util/lib/warning';
 import * as React from 'react';
 import { useImperativeHandle } from 'react';
 import { flushSync } from 'react-dom';
-import isEqual from 'rc-util/lib/isEqual';
-import { getMenuId, IdContext } from './context/IdContext';
+import { IdContext } from './context/IdContext';
 import MenuContextProvider from './context/MenuContext';
 import { PathRegisterContext, PathUserContext } from './context/PathContext';
 import PrivateContext from './context/PrivateContext';
-import useAccessibility from './hooks/useAccessibility';
+import {
+  getFocusableElements,
+  refreshElements,
+  useAccessibility,
+} from './hooks/useAccessibility';
 import useKeyRecords, { OVERFLOW_KEY } from './hooks/useKeyRecords';
 import useMemoCallback from './hooks/useMemoCallback';
 import useUUID from './hooks/useUUID';
@@ -270,8 +274,9 @@ const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
   };
 
   // >>>>> Cache & Reset open keys when inlineCollapsed changed
-  const [inlineCacheOpenKeys, setInlineCacheOpenKeys] =
-    React.useState(mergedOpenKeys);
+  const [inlineCacheOpenKeys, setInlineCacheOpenKeys] = React.useState(
+    mergedOpenKeys,
+  );
 
   const mountRef = React.useRef(false);
 
@@ -347,10 +352,9 @@ const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
     [registerPath, unregisterPath],
   );
 
-  const pathUserContext = React.useMemo(
-    () => ({ isSubPathKey }),
-    [isSubPathKey],
-  );
+  const pathUserContext = React.useMemo(() => ({ isSubPathKey }), [
+    isSubPathKey,
+  ]);
 
   React.useEffect(() => {
     refreshOverflowKeys(
@@ -378,20 +382,31 @@ const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
     setMergedActiveKey(undefined);
   });
 
-  useImperativeHandle(ref, () => ({
-    list: containerRef.current,
-    focus: options => {
-      const shouldFocusKey =
-        mergedActiveKey ?? childList.find(node => !node.props.disabled)?.key;
-      if (shouldFocusKey) {
-        containerRef.current
-          ?.querySelector<HTMLLIElement>(
-            `li[data-menu-id='${getMenuId(uuid, shouldFocusKey as string)}']`,
-          )
-          ?.focus?.(options);
-      }
-    },
-  }));
+  useImperativeHandle(ref, () => {
+    return {
+      list: containerRef.current,
+      focus: options => {
+        const keys = getKeys();
+        const { elements, key2element, element2key } = refreshElements(
+          keys,
+          uuid,
+        );
+        const focusableElements = getFocusableElements(
+          containerRef.current,
+          elements,
+        );
+
+        const shouldFocusKey =
+          mergedActiveKey ?? element2key.get(focusableElements[0]);
+
+        const elementToFocus = key2element.get(shouldFocusKey);
+
+        if (shouldFocusKey && elementToFocus) {
+          elementToFocus?.focus?.(options);
+        }
+      },
+    };
+  });
 
   // ======================== Select ========================
   // >>>>> Select keys

src/hooks/useAccessibility.ts

@@ -1,9 +1,9 @@
-import * as React from 'react';
+import { getFocusNodeList } from 'rc-util/lib/Dom/focus';
 import KeyCode from 'rc-util/lib/KeyCode';
 import raf from 'rc-util/lib/raf';
-import { getFocusNodeList } from 'rc-util/lib/Dom/focus';
-import type { MenuMode } from '../interface';
+import * as React from 'react';
 import { getMenuId } from '../context/IdContext';
+import type { MenuMode } from '../interface';
 
 // destruct to reduce minify size
 const { LEFT, RIGHT, UP, DOWN, ENTER, ESC, HOME, END } = KeyCode;
@@ -134,7 +134,7 @@ function getFocusElement(
 /**
  * Get focusable elements from the element set under provided container
  */
-function getFocusableElements(
+export function getFocusableElements(
   container: HTMLElement,
   elements: Set<HTMLElement>,
 ) {
@@ -181,7 +181,27 @@ function getNextFocusElement(
   return sameLevelFocusableMenuElementList[focusIndex];
 }
 
-export default function useAccessibility<T extends HTMLElement>(
+export const refreshElements = (keys: string[], id: string) => {
+  const elements = new Set<HTMLElement>();
+  const key2element = new Map<string, HTMLElement>();
+  const element2key = new Map<HTMLElement, string>();
+
+  keys.forEach(key => {
+    const element = document.querySelector(
+      `[data-menu-id='${getMenuId(id, key)}']`,
+    ) as HTMLElement;
+
+    if (element) {
+      elements.add(element);
+      element2key.set(element, key);
+      key2element.set(key, element);
+    }
+  });
+
+  return { elements, key2element, element2key };
+};
+
+export function useAccessibility<T extends HTMLElement>(
   mode: MenuMode,
   activeKey: string,
   isRtl: boolean,
@@ -216,35 +236,10 @@ export default function useAccessibility<T extends HTMLElement>(
     const { which } = e;
 
     if ([...ArrowKeys, ENTER, ESC, HOME, END].includes(which)) {
-      // Convert key to elements
-      let elements: Set<HTMLElement>;
-      let key2element: Map<string, HTMLElement>;
-      let element2key: Map<HTMLElement, string>;
-
-      // >>> Wrap as function since we use raf for some case
-      const refreshElements = () => {
-        elements = new Set<HTMLElement>();
-        key2element = new Map();
-        element2key = new Map();
-
-        const keys = getKeys();
-
-        keys.forEach(key => {
-          const element = document.querySelector(
-            `[data-menu-id='${getMenuId(id, key)}']`,
-          ) as HTMLElement;
-
-          if (element) {
-            elements.add(element);
-            element2key.set(element, key);
-            key2element.set(key, element);
-          }
-        });
+      const keys = getKeys();
 
-        return elements;
-      };
-
-      refreshElements();
+      let refreshedElements = refreshElements(keys, id);
+      const { elements, key2element, element2key } = refreshedElements;
 
       // First we should find current focused MenuItem/SubMenu element
       const activeElement = key2element.get(activeKey);
@@ -341,15 +336,15 @@ export default function useAccessibility<T extends HTMLElement>(
         cleanRaf();
         rafRef.current = raf(() => {
           // Async should resync elements
-          refreshElements();
+          refreshedElements = refreshElements(keys, id);
 
           const controlId = focusMenuElement.getAttribute('aria-controls');
           const subQueryContainer = document.getElementById(controlId);
 
           // Get sub focusable menu item
           const targetElement = getNextFocusElement(
             subQueryContainer,
-            elements,
+            refreshedElements.elements,
           );
 
           // Focus menu item

tests/Focus.spec.tsx

@@ -1,9 +1,20 @@
 /* eslint-disable no-undef */
-import { fireEvent, render } from '@testing-library/react';
+import { act, fireEvent, render } from '@testing-library/react';
+import { spyElementPrototypes } from 'rc-util/lib/test/domHook';
 import React from 'react';
-import Menu, { MenuItem, SubMenu } from '../src';
+import Menu, { MenuItem, MenuItemGroup, MenuRef, SubMenu } from '../src';
 
 describe('Focus', () => {
+  beforeAll(() => {
+    // Mock to force make menu item visible
+    spyElementPrototypes(HTMLElement, {
+      offsetParent: {
+        get() {
+          return this.parentElement;
+        },
+      },
+    });
+  });
 
   beforeEach(() => {
     global.triggerProps = null;
@@ -15,13 +26,15 @@ describe('Focus', () => {
     jest.useRealTimers();
   });
 
-  it('Get focus', () => {
-    const { container } = render(
-      <Menu mode="inline" openKeys={['s']}>
-        <SubMenu key="s" title="submenu">
-          <MenuItem key="1">1</MenuItem>
-        </SubMenu>
-      </Menu>,
+  it('Get focus', async () => {
+    const { container } = await act(async () =>
+      render(
+        <Menu mode="inline" openKeys={['s']}>
+          <SubMenu key="s" title="submenu">
+            <MenuItem key="1">1</MenuItem>
+          </SubMenu>
+        </Menu>,
+      ),
     );
 
     // Item focus
@@ -34,5 +47,149 @@ describe('Focus', () => {
     fireEvent.focus(container.querySelector('.rc-menu-submenu-title'));
     expect(container.querySelector('.rc-menu-submenu-active')).toBeTruthy();
   });
+
+  it('should support focus through ref', async () => {
+    const menuRef = React.createRef<MenuRef>();
+    const { getByTestId } = await act(async () =>
+      render(
+        <Menu ref={menuRef}>
+          <SubMenu key="bamboo" title="Disabled" disabled>
+            <MenuItem key="bamboo-child">Disabled child</MenuItem>
+          </SubMenu>
+          <MenuItem key="light" data-testid="first-focusable">
+            Light
+          </MenuItem>
+        </Menu>,
+      ),
+    );
+
+    act(() => menuRef.current.focus());
+
+    const firstFocusableItem = getByTestId('first-focusable');
+    expect(document.activeElement).toBe(firstFocusableItem);
+    expect(firstFocusableItem).toHaveClass('rc-menu-item-active');
+  });
+
+  it('should focus active item through ref', async () => {
+    const menuRef = React.createRef<MenuRef>();
+    const { getByTestId } = await act(async () =>
+      render(
+        <Menu ref={menuRef} activeKey="cat">
+          <MenuItem key="light">Light</MenuItem>
+          <MenuItem key="cat" data-testid="active-key">
+            Cat
+          </MenuItem>
+        </Menu>,
+      ),
+    );
+    act(() => menuRef.current.focus());
+
+    const activeKey = getByTestId('active-key');
+    expect(document.activeElement).toBe(activeKey);
+    expect(activeKey).toHaveClass('rc-menu-item-active');
+  });
+
+  it('focus moves to the next accessible menu item if the first child is empty group', async () => {
+    const menuRef = React.createRef<MenuRef>();
+    const { getByTestId } = await act(async () =>
+      render(
+        <Menu ref={menuRef}>
+          <MenuItemGroup title="group" key="group" />
+          <SubMenu key="bamboo" title="Disabled" disabled>
+            <MenuItem key="bamboo-child">Disabled child</MenuItem>
+          </SubMenu>
+          <MenuItem key="light" data-testid="first-focusable">
+            Light
+          </MenuItem>
+        </Menu>,
+      ),
+    );
+
+    act(() => menuRef.current.focus());
+
+    const firstFocusableItem = getByTestId('first-focusable');
+    expect(document.activeElement).toBe(firstFocusableItem);
+    expect(firstFocusableItem).toHaveClass('rc-menu-item-active');
+  });
+
+  it('focus moves to the next accessible group item if the first child is non-empty group', async () => {
+    const menuRef = React.createRef<MenuRef>();
+    const { getByTestId } = await act(async () =>
+      render(
+        <Menu ref={menuRef}>
+          <MenuItemGroup title="group" key="group">
+            <MenuItem key="group-child-1" disabled>
+              group-child-1
+            </MenuItem>
+            <MenuItem key="group-child-2" data-testid="first-focusable">
+              group-child-2
+            </MenuItem>
+          </MenuItemGroup>
+          <MenuItem key="light">Light</MenuItem>
+        </Menu>,
+      ),
+    );
+
+    act(() => menuRef.current.focus());
+
+    const firstFocusableItem = getByTestId('first-focusable');
+    expect(document.activeElement).toBe(firstFocusableItem);
+    expect(firstFocusableItem).toHaveClass('rc-menu-item-active');
+  });
+
+  it('focus moves to nested group item correctly', async () => {
+    const menuRef = React.createRef<MenuRef>();
+    const { getByTestId } = await act(async () =>
+      render(
+        <Menu ref={menuRef}>
+          <MenuItemGroup title="group" key="group">
+            <MenuItem key="group-child-1" disabled>
+              group-child-1
+            </MenuItem>
+            <MenuItemGroup title="nested group" key="nested-group">
+              <MenuItem key="nested-group-child-1" disabled>
+                nested-group-child-1
+              </MenuItem>
+              <MenuItem
+                key="nested-group-child-2"
+                data-testid="first-focusable"
+              >
+                nested-group-child-2
+              </MenuItem>
+            </MenuItemGroup>
+            <MenuItem key="group-child-3">group-child-3</MenuItem>
+          </MenuItemGroup>
+        </Menu>,
+      ),
+    );
+
+    act(() => menuRef.current.focus());
+
+    const firstFocusableItem = getByTestId('first-focusable');
+    expect(document.activeElement).toBe(firstFocusableItem);
+    expect(firstFocusableItem).toHaveClass('rc-menu-item-active');
+  });
+
+  it('focus moves to submenu correctly', async () => {
+    const menuRef = React.createRef<MenuRef>();
+    const { getByTestId, getByTitle } = await act(async () =>
+      render(
+        <Menu ref={menuRef}>
+          <SubMenu key="sub-menu-disabled" title="Disabled" disabled>
+            <MenuItem key="sub-menu-disabled-child">Disabled child</MenuItem>
+          </SubMenu>
+          <SubMenu key="sub-menu" data-testid="sub-menu" title="Submenu">
+            <MenuItem key="sub-menu-child-1">Submenu child</MenuItem>
+          </SubMenu>
+          <MenuItem key="light">Light</MenuItem>
+        </Menu>,
+      ),
+    );
+
+    act(() => menuRef.current.focus());
+
+    expect(document.activeElement).toBe(getByTitle('Submenu'));
+    expect(getByTestId('sub-menu')).toHaveClass('rc-menu-submenu-active');
+  });
 });
 /* eslint-enable */