Fix refocus button bug

CHANGELOG.md

@@ -9,8 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixes
 
-- Fixed `outside click` not re-focusing the `Menu.Button` ([#220](https://github.com/tailwindlabs/headlessui/pull/220))
-- Fixed `outside click` not re-focusing the `Listbox.Button` ([#220](https://github.com/tailwindlabs/headlessui/pull/220))
+- Fixed `outside click` not re-focusing the `Menu.Button` ([#220](https://github.com/tailwindlabs/headlessui/pull/220), [#256](https://github.com/tailwindlabs/headlessui/pull/256))
+- Fixed `outside click` not re-focusing the `Listbox.Button` ([#220](https://github.com/tailwindlabs/headlessui/pull/220), [#256](https://github.com/tailwindlabs/headlessui/pull/256))
 - Fix incorrect type error `unique symbol` ([#248](https://github.com/tailwindlabs/headlessui/pull/248), [#240](https://github.com/tailwindlabs/headlessui/issues/240))
 
 ### Added

packages/@headlessui-react/src/components/listbox/listbox.test.tsx

@@ -34,6 +34,7 @@ import {
   getListboxOptions,
   getListboxLabel,
   ListboxState,
+  getByText,
 } from '../../test-utils/accessibility-assertions'
 
 jest.mock('../../hooks/use-id')
@@ -2878,7 +2879,7 @@ describe('Mouse interactions', () => {
       })
       assertListbox({ state: ListboxState.InvisibleUnmounted })
 
-      // Try to open the menu
+      // Try to open the listbox
       await click(getListboxButton(), MouseButton.Right)
 
       // Verify it is still closed
@@ -3071,19 +3072,19 @@ describe('Mouse interactions', () => {
 
       let [button1, button2] = getListboxButtons()
 
-      // Click the first menu button
+      // Click the first listbox button
       await click(button1)
-      expect(getListboxes()).toHaveLength(1) // Only 1 menu should be visible
+      expect(getListboxes()).toHaveLength(1) // Only 1 listbox should be visible
 
-      // Ensure the open menu is linked to the first button
+      // Ensure the open listbox is linked to the first button
       assertListboxButtonLinkedWithListbox(button1, getListbox())
 
-      // Click the second menu button
+      // Click the second listbox button
       await click(button2)
 
-      expect(getListboxes()).toHaveLength(1) // Only 1 menu should be visible
+      expect(getListboxes()).toHaveLength(1) // Only 1 listbox should be visible
 
-      // Ensure the open menu is linked to the second button
+      // Ensure the open listbox is linked to the second button
       assertListboxButtonLinkedWithListbox(button2, getListbox())
     })
   )
@@ -3118,6 +3119,47 @@ describe('Mouse interactions', () => {
     })
   )
 
+  it(
+    'should be possible to click outside of the listbox, on an element which is within a focusable element, which closes the listbox',
+    suppressConsoleLogs(async () => {
+      let focusFn = jest.fn()
+      render(
+        <div>
+          <Listbox value={undefined} onChange={console.log}>
+            <Listbox.Button onFocus={focusFn}>Trigger</Listbox.Button>
+            <Listbox.Options>
+              <Listbox.Option value="alice">alice</Listbox.Option>
+              <Listbox.Option value="bob">bob</Listbox.Option>
+              <Listbox.Option value="charlie">charlie</Listbox.Option>
+            </Listbox.Options>
+          </Listbox>
+
+          <button id="btn">
+            <span>Next</span>
+          </button>
+        </div>
+      )
+
+      // Click the listbox button
+      await click(getListboxButton())
+
+      // Ensure the listbox is open
+      assertListbox({ state: ListboxState.Visible })
+
+      // Click the span inside the button
+      await click(getByText('Next'))
+
+      // Ensure the listbox is closed
+      assertListbox({ state: ListboxState.InvisibleUnmounted })
+
+      // Ensure the outside button is focused
+      assertActiveElement(document.getElementById('btn'))
+
+      // Ensure that the focus button only got focus once (first click)
+      expect(focusFn).toHaveBeenCalledTimes(1)
+    })
+  )
+
   it(
     'should be possible to hover an option and make it active',
     suppressConsoleLogs(async () => {

packages/@headlessui-react/src/components/listbox/listbox.tsx

@@ -31,7 +31,7 @@ import { Keys } from '../keyboard'
 import { Focus, calculateActiveIndex } from '../../utils/calculate-active-index'
 import { resolvePropValue } from '../../utils/resolve-prop-value'
 import { isDisabledReactIssue7711 } from '../../utils/bugs'
-import { isFocusableElement } from '../../utils/focus-management'
+import { isFocusableElement, FocusableMode } from '../../utils/focus-management'
 
 enum ListboxStates {
   Open,
@@ -229,7 +229,7 @@ export function Listbox<TTag extends ElementType = typeof DEFAULT_LISTBOX_TAG, T
 
       dispatch({ type: ActionTypes.CloseListbox })
 
-      if (!isFocusableElement(target)) {
+      if (!isFocusableElement(target, FocusableMode.Loose)) {
         event.preventDefault()
         buttonRef.current?.focus()
       }

packages/@headlessui-react/src/components/menu/menu.test.tsx

@@ -17,6 +17,7 @@ import {
   getMenu,
   getMenus,
   getMenuItems,
+  getByText,
 } from '../../test-utils/accessibility-assertions'
 import {
   click,
@@ -2665,6 +2666,47 @@ describe('Mouse interactions', () => {
     })
   )
 
+  it(
+    'should be possible to click outside of the menu, on an element which is within a focusable element, which closes the menu',
+    suppressConsoleLogs(async () => {
+      let focusFn = jest.fn()
+      render(
+        <div>
+          <Menu>
+            <Menu.Button onFocus={focusFn}>Trigger</Menu.Button>
+            <Menu.Items>
+              <Menu.Item as="a">alice</Menu.Item>
+              <Menu.Item as="a">bob</Menu.Item>
+              <Menu.Item as="a">charlie</Menu.Item>
+            </Menu.Items>
+          </Menu>
+
+          <button id="btn">
+            <span>Next</span>
+          </button>
+        </div>
+      )
+
+      // Click the menu button
+      await click(getMenuButton())
+
+      // Ensure the menu is open
+      assertMenu({ state: MenuState.Visible })
+
+      // Click the span inside the button
+      await click(getByText('Next'))
+
+      // Ensure the menu is closed
+      assertMenu({ state: MenuState.InvisibleUnmounted })
+
+      // Ensure the outside button is focused
+      assertActiveElement(document.getElementById('btn'))
+
+      // Ensure that the focus button only got focus once (first click)
+      expect(focusFn).toHaveBeenCalledTimes(1)
+    })
+  )
+
   it(
     'should be possible to hover an item and make it active',
     suppressConsoleLogs(async () => {

packages/@headlessui-react/src/components/menu/menu.tsx

@@ -31,7 +31,7 @@ import { Keys } from '../keyboard'
 import { Focus, calculateActiveIndex } from '../../utils/calculate-active-index'
 import { resolvePropValue } from '../../utils/resolve-prop-value'
 import { isDisabledReactIssue7711 } from '../../utils/bugs'
-import { isFocusableElement } from '../../utils/focus-management'
+import { isFocusableElement, FocusableMode } from '../../utils/focus-management'
 
 enum MenuStates {
   Open,
@@ -185,7 +185,7 @@ export function Menu<TTag extends ElementType = typeof DEFAULT_MENU_TAG>(
 
       dispatch({ type: ActionTypes.CloseMenu })
 
-      if (!isFocusableElement(target)) {
+      if (!isFocusableElement(target, FocusableMode.Loose)) {
         event.preventDefault()
         buttonRef.current?.focus()
       }

packages/@headlessui-react/src/test-utils/interactions.ts

@@ -191,9 +191,14 @@ export async function click(
         fireEvent.mouseDown(element, options)
       }
 
-      // Ensure to trigger a `focus` event if the element is focusable
-      if ((element as HTMLElement)?.matches(focusableSelector)) {
-        ;(element as HTMLElement).focus()
+      // Ensure to trigger a `focus` event if the element is focusable, or within a focusable element
+      let next: HTMLElement | null = element as HTMLElement | null
+      while (next !== null) {
+        if (next.matches(focusableSelector)) {
+          next.focus()
+          break
+        }
+        next = next.parentElement
       }
 
       fireEvent.pointerUp(element, options)

packages/@headlessui-react/src/utils/focus-management.ts

@@ -1,3 +1,5 @@
+import { match } from './match'
+
 // Credit:
 //  - https://stackoverflow.com/a/30753870
 let focusableSelector = [
@@ -22,22 +24,22 @@ let focusableSelector = [
   .join(',')
 
 export enum Focus {
-  /* Focus the first non-disabled element */
+  /** Focus the first non-disabled element */
   First = 1 << 0,
 
-  /* Focus the previous non-disabled element */
+  /** Focus the previous non-disabled element */
   Previous = 1 << 1,
 
-  /* Focus the next non-disabled element */
+  /** Focus the next non-disabled element */
   Next = 1 << 2,
 
-  /* Focus the last non-disabled element */
+  /** Focus the last non-disabled element */
   Last = 1 << 3,
 
-  /* Wrap tab around */
+  /** Wrap tab around */
   WrapAround = 1 << 4,
 
-  /* Prevent scrolling the focusable elements into view */
+  /** Prevent scrolling the focusable elements into view */
   NoScroll = 1 << 5,
 }
 
@@ -58,8 +60,35 @@ export function getFocusableElements(container: HTMLElement | null = document.bo
   return Array.from(container.querySelectorAll<HTMLElement>(focusableSelector))
 }
 
-export function isFocusableElement(element: HTMLElement) {
-  return element.matches(focusableSelector)
+export enum FocusableMode {
+  /** The element itself must be focusable. */
+  Strict,
+
+  /** The element should be inside of a focusable element. */
+  Loose,
+}
+
+export function isFocusableElement(
+  element: HTMLElement,
+  mode: FocusableMode = FocusableMode.Strict
+) {
+  if (element === document.body) return false
+
+  return match(mode, {
+    [FocusableMode.Strict]() {
+      return element.matches(focusableSelector)
+    },
+    [FocusableMode.Loose]() {
+      let next: HTMLElement | null = element
+
+      while (next !== null) {
+        if (next.matches(focusableSelector)) return true
+        next = next.parentElement
+      }
+
+      return false
+    },
+  })
 }
 
 export function focusElement(element: HTMLElement | null) {