Export useFocusTrap hook

.changeset/tidy-traps-focus.md

@@ -0,0 +1,5 @@
+---
+'@primer/react-brand': minor
+---
+
+Exported `useFocusTrap` from the package root and created stable container and initial-focus refs when they are not provided.

packages/react/src/hooks/useFocusTrap.ssr.test.tsx

@@ -0,0 +1,16 @@
+/** @jest-environment node */
+
+import React from 'react'
+import {renderToString} from 'react-dom/server'
+import {useFocusTrap} from './useFocusTrap'
+
+describe('useFocusTrap SSR', () => {
+  it('renders without a document', () => {
+    const TestComponent = () => {
+      useFocusTrap()
+      return null
+    }
+
+    expect(() => renderToString(<TestComponent />)).not.toThrow()
+  })
+})

packages/react/src/hooks/useFocusTrap.test.tsx

@@ -0,0 +1,112 @@
+import React from 'react'
+import {focusTrap} from '@primer/behaviors'
+import {render, renderHook} from '@testing-library/react'
+import {useFocusTrap} from './useFocusTrap'
+
+jest.mock('@primer/behaviors', () => ({
+  focusTrap: jest.fn(),
+}))
+
+const mockFocusTrap = jest.mocked(focusTrap)
+
+describe('useFocusTrap', () => {
+  beforeEach(() => {
+    mockFocusTrap.mockReturnValue(new AbortController())
+  })
+
+  afterEach(() => {
+    jest.clearAllMocks()
+  })
+
+  it('creates stable refs when refs are not provided', () => {
+    const {result, rerender} = renderHook(() => useFocusTrap({disabled: true}))
+    const initialContainerRef = result.current.containerRef
+    const initialFocusRef = result.current.initialFocusRef
+
+    expect(initialContainerRef.current).toBeNull()
+    expect(initialFocusRef.current).toBeNull()
+
+    rerender()
+
+    expect(result.current.containerRef).toBe(initialContainerRef)
+    expect(result.current.initialFocusRef).toBe(initialFocusRef)
+  })
+
+  it('preserves provided refs', () => {
+    const containerRef = React.createRef<HTMLDivElement>()
+    const initialFocusRef = React.createRef<HTMLButtonElement>()
+    const {result} = renderHook(() => useFocusTrap({containerRef, initialFocusRef, disabled: true}))
+
+    expect(result.current.containerRef).toBe(containerRef)
+    expect(result.current.initialFocusRef).toBe(initialFocusRef)
+  })
+
+  it('passes generated ref elements to the focus trap', () => {
+    const TestComponent = () => {
+      const {containerRef, initialFocusRef} = useFocusTrap<HTMLDivElement, HTMLButtonElement>()
+
+      return (
+        <div ref={containerRef}>
+          <button ref={initialFocusRef} />
+        </div>
+      )
+    }
+
+    const {container} = render(<TestComponent />)
+
+    expect(mockFocusTrap).toHaveBeenCalledWith(container.querySelector('div'), container.querySelector('button'))
+  })
+
+  it('starts and aborts the focus trap when disabled changes', () => {
+    const abortController = new AbortController()
+    const abortSpy = jest.spyOn(abortController, 'abort')
+    mockFocusTrap.mockReturnValue(abortController)
+
+    const TestComponent = ({disabled}: {disabled: boolean}) => {
+      const {containerRef} = useFocusTrap<HTMLDivElement>({disabled})
+      return <div ref={containerRef} />
+    }
+
+    const {rerender} = render(<TestComponent disabled />)
+
+    expect(mockFocusTrap).not.toHaveBeenCalled()
+
+    rerender(<TestComponent disabled={false} />)
+
+    expect(mockFocusTrap).toHaveBeenCalledTimes(1)
+
+    rerender(<TestComponent disabled />)
+
+    expect(abortSpy).toHaveBeenCalled()
+  })
+
+  it('captures fresh focus after cleaning up a non-HTMLElement active element', () => {
+    const svg = document.createElementNS('http://www.w3.org/2000/svg', 'svg')
+    svg.setAttribute('tabindex', '0')
+    document.body.append(svg)
+    svg.focus()
+
+    const nextFocusedElement = document.createElement('button')
+    document.body.append(nextFocusedElement)
+
+    const TestComponent = ({disabled}: {disabled: boolean}) => {
+      const {containerRef} = useFocusTrap<HTMLDivElement>({disabled, restoreFocusOnCleanUp: true})
+      return <div ref={containerRef} />
+    }
+
+    const {rerender} = render(<TestComponent disabled={false} />)
+
+    rerender(<TestComponent disabled />)
+    nextFocusedElement.focus()
+
+    const focusSpy = jest.spyOn(nextFocusedElement, 'focus')
+
+    rerender(<TestComponent disabled={false} />)
+    rerender(<TestComponent disabled />)
+
+    expect(focusSpy).toHaveBeenCalled()
+
+    svg.remove()
+    nextFocusedElement.remove()
+  })
+})

packages/react/src/hooks/useFocusTrap.ts

@@ -1,18 +1,22 @@
 import React from 'react'
 import {focusTrap} from '@primer/behaviors'
+import {useProvidedRefOrCreate} from './useRef'
 
-export interface FocusTrapHookSettings {
+export interface FocusTrapHookSettings<
+  ContainerElement extends HTMLElement = HTMLElement,
+  InitialFocusElement extends HTMLElement = HTMLElement,
+> {
   /**
-   * Ref that will be used for the trapping container. If not provided, one will
+   * Ref object that will be used for the trapping container. If not provided, one will
    * be created by this hook and returned.
    */
-  containerRef: React.RefObject<HTMLElement | null>
+  containerRef?: React.RefObject<ContainerElement | null>
 
   /**
-   * Ref for the element that should receive focus when the focus trap is first enabled.
+   * Ref object for the element that should receive focus when the focus trap is first enabled.
    * If not provided, one will be created by this hook and returned. Its use is optional.
    */
-  initialFocusRef?: React.RefObject<HTMLElement | null>
+  initialFocusRef?: React.RefObject<InitialFocusElement | null>
 
   /**
    * Set to true to disable the focus trap and clean up listeners. Can be re-enabled at any time.
@@ -30,40 +34,40 @@ export interface FocusTrapHookSettings {
  * that should trap focus.
  * @param settings {FocusTrapHookSettings}
  */
-export function useFocusTrap(
-  settings?: FocusTrapHookSettings,
+export function useFocusTrap<
+  ContainerElement extends HTMLElement = HTMLElement,
+  InitialFocusElement extends HTMLElement = HTMLElement,
+>(
+  settings?: FocusTrapHookSettings<ContainerElement, InitialFocusElement>,
   dependencies: React.DependencyList = [],
 ): {
-  containerRef: React.RefObject<HTMLElement | null> | undefined
-  initialFocusRef: React.RefObject<HTMLElement | null> | undefined
+  containerRef: React.RefObject<ContainerElement | null>
+  initialFocusRef: React.RefObject<InitialFocusElement | null>
 } {
-  const containerRef = settings?.containerRef
-  const initialFocusRef = settings?.initialFocusRef
+  const containerRef = useProvidedRefOrCreate<ContainerElement | null>(settings?.containerRef)
+  const initialFocusRef = useProvidedRefOrCreate<InitialFocusElement | null>(settings?.initialFocusRef)
   const disabled = settings?.disabled
   const abortController = React.useRef<AbortController | null>(null)
   const previousFocusedElement = React.useRef<Element | null>(null)
 
-  // If we are enabling a focus trap and haven't already stored the previously focused element
-  // go ahead an do that so we can restore later when the trap is disabled.
-  if (!previousFocusedElement.current && !settings?.disabled) {
-    previousFocusedElement.current = document.activeElement
-  }
-
   // This function removes the event listeners that enable the focus trap and restores focus
   // to the previously-focused element (if necessary).
   function disableTrap() {
     abortController.current?.abort()
     if (settings?.restoreFocusOnCleanUp && previousFocusedElement.current instanceof HTMLElement) {
       previousFocusedElement.current.focus()
-      previousFocusedElement.current = null
     }
+    previousFocusedElement.current = null
   }
 
   React.useEffect(
     () => {
-      if (containerRef?.current instanceof HTMLElement) {
+      if (containerRef.current instanceof HTMLElement) {
         if (!disabled) {
-          abortController.current = focusTrap(containerRef.current, initialFocusRef?.current ?? undefined) ?? null
+          if (!previousFocusedElement.current) {
+            previousFocusedElement.current = document.activeElement
+          }
+          abortController.current = focusTrap(containerRef.current, initialFocusRef.current ?? undefined) ?? null
           return () => {
             disableTrap()
           }

packages/react/src/index.ts

@@ -59,4 +59,5 @@ export * from './TextCursorAnimation'
 export * from './Tiles'
 
 // hooks
+export * from './hooks/useFocusTrap'
 export * from './hooks/useWindowSize'