ActionMenu2 + DropdownMenu2: Initial focus is based on key used to open menu

.changeset/menus-initial-focus.md

@@ -0,0 +1,8 @@
+---
+'@primer/react': patch
+---
+
+ActionMenu2 + DropdownMenu2: Focus the correct element when Menu is opened with keyboard. [See detailed spec.](https://github.com/github/primer/issues/518#issuecomment-999104848)
+
+- ArrowDown | Space | Enter: first element
+- ArrowUp: last element

src/ActionMenu2.tsx

@@ -3,7 +3,7 @@ import {useSSRSafeId} from '@react-aria/ssr'
 import {TriangleDownIcon} from '@primer/octicons-react'
 import {AnchoredOverlay, AnchoredOverlayProps} from './AnchoredOverlay'
 import {OverlayProps} from './Overlay'
-import {useProvidedRefOrCreate, useProvidedStateOrCreate} from './hooks'
+import {useProvidedRefOrCreate, useProvidedStateOrCreate, useMenuInitialFocus} from './hooks'
 import {Divider} from './ActionList2/Divider'
 import {ActionListContainerContext} from './ActionList2/ActionListContainerContext'
 import {Button, ButtonProps} from './Button2'
@@ -95,27 +95,31 @@ const Overlay: React.FC<MenuOverlayProps> = ({children, ...overlayProps}) => {
     'anchorRef'
   >
 
+  const {containerRef, openWithFocus} = useMenuInitialFocus(open, onOpen)
+
   return (
     <AnchoredOverlay
       anchorRef={anchorRef}
       renderAnchor={renderAnchor}
       anchorId={anchorId}
       open={open}
-      onOpen={onOpen}
+      onOpen={openWithFocus}
       onClose={onClose}
       overlayProps={overlayProps}
     >
-      <ActionListContainerContext.Provider
-        value={{
-          container: 'ActionMenu',
-          listRole: 'menu',
-          itemRole: 'menuitem',
-          listLabelledBy: anchorId,
-          afterSelect: onClose
-        }}
-      >
-        {children}
-      </ActionListContainerContext.Provider>
+      <div ref={containerRef}>
+        <ActionListContainerContext.Provider
+          value={{
+            container: 'ActionMenu',
+            listRole: 'menu',
+            itemRole: 'menuitem',
+            listLabelledBy: anchorId,
+            afterSelect: onClose
+          }}
+        >
+          {children}
+        </ActionListContainerContext.Provider>
+      </div>
     </AnchoredOverlay>
   )
 }

src/AnchoredOverlay/AnchoredOverlay.tsx

@@ -55,7 +55,7 @@ interface AnchoredOverlayBaseProps extends Pick<OverlayProps, 'height' | 'width'
   /**
    * A callback which is called whenever the overlay is currently closed and an "open gesture" is detected.
    */
-  onOpen?: (gesture: 'anchor-click' | 'anchor-key-press') => unknown
+  onOpen?: (gesture: 'anchor-click' | 'anchor-key-press', event?: React.KeyboardEvent<HTMLElement>) => unknown
 
   /**
    * A callback which is called whenever the overlay is currently open and a "close gesture" is detected.
@@ -113,7 +113,7 @@ export const AnchoredOverlay: React.FC<AnchoredOverlayProps> = ({
     (event: React.KeyboardEvent<HTMLElement>) => {
       if (!event.defaultPrevented) {
         if (!open && ['ArrowDown', 'ArrowUp', ' ', 'Enter'].includes(event.key)) {
-          onOpen?.('anchor-key-press')
+          onOpen?.('anchor-key-press', event)
           event.preventDefault()
         }
       }

src/DropdownMenu2.tsx

@@ -4,7 +4,7 @@ import {TriangleDownIcon} from '@primer/octicons-react'
 import {Button, ButtonProps} from './Button2'
 import {AnchoredOverlay, AnchoredOverlayProps} from './AnchoredOverlay'
 import {OverlayProps} from './Overlay'
-import {useProvidedRefOrCreate, useProvidedStateOrCreate} from './hooks'
+import {useProvidedRefOrCreate, useProvidedStateOrCreate, useMenuInitialFocus} from './hooks'
 import {Divider} from './ActionList2/Divider'
 import {ActionListContainerContext} from './ActionList2/ActionListContainerContext'
 import {MandateProps} from './utils/types'
@@ -99,29 +99,33 @@ const Overlay: React.FC<MenuOverlayProps> = ({children, ...overlayProps}) => {
     'anchorRef'
   >
 
+  const {containerRef, openWithFocus} = useMenuInitialFocus(open, onOpen)
+
   return (
     <AnchoredOverlay
       anchorRef={anchorRef}
       renderAnchor={renderAnchor}
       anchorId={anchorId}
       open={open}
-      onOpen={onOpen}
+      onOpen={openWithFocus}
       onClose={onClose}
       overlayProps={overlayProps}
     >
-      <ActionListContainerContext.Provider
-        value={{
-          container: 'DropdownMenu',
-          listRole: 'menu',
-          itemRole: 'menuitemradio',
-          listLabelledBy: anchorId,
-          selectionVariant: 'single',
-          selectionAttribute: 'aria-checked',
-          afterSelect: onClose
-        }}
-      >
-        {children}
-      </ActionListContainerContext.Provider>
+      <div ref={containerRef}>
+        <ActionListContainerContext.Provider
+          value={{
+            container: 'DropdownMenu',
+            listRole: 'menu',
+            itemRole: 'menuitemradio',
+            listLabelledBy: anchorId,
+            selectionVariant: 'single',
+            selectionAttribute: 'aria-checked',
+            afterSelect: onClose
+          }}
+        >
+          {children}
+        </ActionListContainerContext.Provider>
+      </div>
     </AnchoredOverlay>
   )
 }

src/__tests__/hooks/useMenuInitialFocus.test.tsx

@@ -0,0 +1,84 @@
+import React from 'react'
+import {render, fireEvent, cleanup, waitFor} from '@testing-library/react'
+import {useMenuInitialFocus} from '../../hooks'
+
+const Component = () => {
+  const [open, setOpen] = React.useState(false)
+  const onOpen = () => setOpen(!open)
+  const {containerRef, openWithFocus} = useMenuInitialFocus(open, onOpen)
+
+  return (
+    <>
+      <button onClick={() => setOpen(true)} onKeyDown={event => openWithFocus('anchor-key-press', event)}>
+        open container
+      </button>
+      {open && (
+        <div ref={containerRef}>
+          <span>not focusable</span>
+          <button>first focusable element</button>
+          <button>second focusable element</button>
+          <button>third focusable element</button>
+          <span>not focusable</span>
+        </div>
+      )}
+    </>
+  )
+}
+
+describe('useMenuFocus', () => {
+  afterEach(cleanup)
+
+  it('should focus first element when opened with Enter', async () => {
+    const {getByText} = render(<Component />)
+    const button = getByText('open container')
+
+    fireEvent.keyDown(button, {key: 'Enter', code: 'Enter', keyCode: 13, charCode: 13})
+
+    /** We use waitFor because the hook uses an effect with setTimeout
+     *  and we need to wait for that to happen in the next tick
+     */
+    await waitFor(() => {
+      const firstButton = getByText('first focusable element')
+      expect(firstButton).toEqual(document.activeElement)
+    })
+  })
+
+  it('should focus first element when opened with ArrowDown', async () => {
+    const {getByText} = render(<Component />)
+    const button = getByText('open container')
+
+    fireEvent.keyDown(button, {key: 'ArrowDown', code: 'ArrowDown', keyCode: 40, charCode: 40})
+
+    await waitFor(() => {
+      const firstButton = getByText('first focusable element')
+      expect(firstButton).toEqual(document.activeElement)
+    })
+  })
+
+  it('should focus last element when opened with ArrowUp', async () => {
+    const {getByText} = render(<Component />)
+    const button = getByText('open container')
+
+    fireEvent.keyDown(button, {key: 'ArrowUp', code: 'ArrowUp', keyCode: 38, charCode: 38})
+
+    await waitFor(() => {
+      const thirdButton = getByText('third focusable element')
+      expect(thirdButton).toEqual(document.activeElement)
+    })
+  })
+
+  it('should focus neither when a different letter is pressed', async () => {
+    const {getByText} = render(<Component />)
+    const button = getByText('open container')
+
+    fireEvent.keyDown(button, {key: 'ArrowRight', code: 'ArrowRight', keyCode: 39, charCode: 39})
+
+    await waitFor(() => {
+      const firstButton = getByText('first focusable element')
+      const thirdButton = getByText('third focusable element')
+      expect(firstButton).not.toEqual(document.activeElement)
+      expect(thirdButton).not.toEqual(document.activeElement)
+      expect(document.body).toEqual(document.activeElement)
+    })
+  })
+})

src/hooks/index.ts

@@ -10,3 +10,4 @@ export {useOverlay} from './useOverlay'
 export type {UseOverlaySettings} from './useOverlay'
 export {useRenderForcingRef} from './useRenderForcingRef'
 export {useProvidedStateOrCreate} from './useProvidedStateOrCreate'
+export {useMenuInitialFocus} from './useMenuInitialFocus'

src/hooks/useMenuInitialFocus.ts

@@ -0,0 +1,40 @@
+import React from 'react'
+import {iterateFocusableElements} from '@primer/behaviors/utils'
+
+type Gesture = 'anchor-click' | 'anchor-key-press'
+type Callback = (gesture: Gesture, event?: React.KeyboardEvent<HTMLElement>) => unknown
+
+export const useMenuInitialFocus = (open: boolean, onOpen?: Callback) => {
+  const containerRef = React.createRef<HTMLDivElement>()
+  const [openingKey, setOpeningKey] = React.useState<string | undefined>(undefined)
+
+  const openWithFocus: Callback = (gesture, event) => {
+    if (gesture === 'anchor-key-press' && event) setOpeningKey(event.code)
+    else setOpeningKey(undefined)
+    if (typeof onOpen === 'function') onOpen(gesture, event)
+  }
+
+  /**
+   * Pick the first element to focus based on the key used to open the Menu
+   * ArrowDown | Space | Enter: first element
+   * ArrowUp: last element
+   */
+
+  React.useEffect(() => {
+    if (!open) return
+    if (!openingKey || !containerRef.current) return
+
+    const iterable = iterateFocusableElements(containerRef.current)
+    if (['ArrowDown', 'Space', 'Enter'].includes(openingKey)) {
+      const firstElement = iterable.next().value
+      /** We push imperative focus to the next tick to prevent React's batching */
+      setTimeout(() => firstElement?.focus())
+    } else if (['ArrowUp'].includes(openingKey)) {
+      const elements = [...iterable]
+      const lastElement = elements[elements.length - 1]
+      setTimeout(() => lastElement.focus())
+    }
+  }, [open, openingKey, containerRef])
+
+  return {containerRef, openWithFocus}
+}