Fix Select keyboard focus trap

packages/@react-aria/select/package.json

@@ -22,6 +22,7 @@
     "url": "https://github.com/adobe/react-spectrum"
   },
   "dependencies": {
+    "@react-aria/focus": "^3.18.3",
     "@react-aria/form": "^3.0.9",
     "@react-aria/i18n": "^3.12.3",
     "@react-aria/interactions": "^3.22.3",

packages/@react-aria/select/src/HiddenSelect.tsx

@@ -11,6 +11,7 @@
  */
 
 import {FocusableElement, RefObject} from '@react-types/shared';
+import {focusSafely, getFocusableTreeWalker} from '@react-aria/focus';
 import React, {ReactNode, useRef} from 'react';
 import {selectData} from './useSelect';
 import {SelectState} from '@react-stately/select';
@@ -108,7 +109,31 @@ export function useHiddenSelect<T>(props: AriaHiddenSelectOptions, state: Select
       type: 'text',
       tabIndex: modality == null || state.isFocused || state.isOpen ? -1 : 0,
       style: {fontSize: 16},
-      onFocus: () => triggerRef.current.focus(),
+      onFocus: (e) => {
+        let position = e.relatedTarget?.compareDocumentPosition(e.currentTarget);
+        let walker = getFocusableTreeWalker(document.body, {tabbable: true});
+        if (position & Node.DOCUMENT_POSITION_PRECEDING) {
+          walker.currentNode = e.relatedTarget;
+          // if focus is coming from "after" this element, we know it's a shift tab
+          let prevNode = walker.previousNode() as FocusableElement | null;
+          if (prevNode) {
+            focusElement(prevNode, true);
+          } else {
+            triggerRef.current.focus();
+          }
+        } else if (position & Node.DOCUMENT_POSITION_FOLLOWING) {
+          walker.currentNode = e.relatedTarget;
+          // if focus is coming from "before" this element, we know it's a tab
+          let nextNode = walker.nextNode() as FocusableElement | null;
+          if (nextNode) {
+            focusElement(nextNode, true);
+          } else {
+            triggerRef.current.focus();
+          }
+        } else {
+          triggerRef.current.focus();
+        }
+      },
       disabled: isDisabled
     },
     selectProps: {
@@ -172,3 +197,19 @@ export function HiddenSelect<T>(props: HiddenSelectProps<T>) {
 
   return null;
 }
+
+function focusElement(element: FocusableElement | null, scroll = false) {
+  if (element != null && !scroll) {
+    try {
+      focusSafely(element);
+    } catch (err) {
+      // ignore
+    }
+  } else if (element != null) {
+    try {
+      element.focus();
+    } catch (err) {
+      // ignore
+    }
+  }
+}

packages/react-aria-components/test/Select.test.js

@@ -11,7 +11,7 @@
  */
 
 import {act, pointerMap, render, within} from '@react-spectrum/test-utils-internal';
-import {Button, FieldError, Label, ListBox, ListBoxItem, Popover, Select, SelectContext, SelectValue, Text} from '../';
+import {Button, FieldError, Label, ListBox, ListBoxItem, Popover, Select, SelectContext, SelectStateContext, SelectValue, Text} from '../';
 import React from 'react';
 import {User} from '@react-aria/test-utils';
 import userEvent from '@testing-library/user-event';
@@ -278,4 +278,78 @@ describe('Select', () => {
 
     expect(button).toHaveTextContent('0');
   });
+
+  function SelectClearButton() {
+    let state = React.useContext(SelectStateContext);
+    return (
+      <Button
+        data-testid="clear"
+        // Don't inherit behavior from Select.
+        slot={null}
+        style={{fontSize: 'small', marginTop: 6, padding: 4}}
+        onPress={() => state?.setSelectedKey(null)}>
+        Clear
+      </Button>
+    );
+  }
+
+  it('should support extra children for use with the state', async () => {
+    let onChangeSpy = jest.fn();
+    let {getByTestId} = render(
+      <>
+        <input data-testid="before" />
+        <Select onSelectionChange={onChangeSpy}>
+          <Label>Favorite Animal</Label>
+          <Button data-testid="select">
+            <SelectValue />
+            <span aria-hidden="true">▼</span>
+          </Button>
+          <SelectClearButton />
+          <Popover>
+            <ListBox>
+              <ListBoxItem id="cat">Cat</ListBoxItem>
+              <ListBoxItem id="dog">Dog</ListBoxItem>
+              <ListBoxItem id="kangaroo">Kangaroo</ListBoxItem>
+            </ListBox>
+          </Popover>
+        </Select>
+        <input data-testid="after" />
+      </>
+    );
+
+    let beforeInput = getByTestId('before');
+    let afterInput = getByTestId('after');
+    let wrapper = getByTestId('select');
+    let clearButton = getByTestId('clear');
+    let selectTester = testUtilUser.createTester('Select', {root: wrapper});
+
+    await user.tab();
+    await user.tab();
+    expect(document.activeElement).toBe(selectTester.trigger);
+
+    await user.tab();
+    expect(document.activeElement).toBe(clearButton);
+
+    await user.tab();
+    expect(document.activeElement).toBe(afterInput);
+
+    await user.tab({shift: true});
+    expect(document.activeElement).toBe(clearButton);
+
+    await user.tab({shift: true});
+    expect(document.activeElement).toBe(selectTester.trigger);
+
+    await user.tab({shift: true});
+    expect(document.activeElement).toBe(beforeInput);
+
+    await user.tab();
+    await selectTester.selectOption({optionText: 'Dog'});
+
+    expect(onChangeSpy).toHaveBeenCalledTimes(1);
+    expect(onChangeSpy).toHaveBeenLastCalledWith('dog');
+
+    await user.click(clearButton);
+    expect(onChangeSpy).toHaveBeenCalledTimes(2);
+    expect(onChangeSpy).toHaveBeenLastCalledWith(null);
+  });
 });

yarn.lock

@@ -6171,6 +6171,7 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "@react-aria/select@workspace:packages/@react-aria/select"
   dependencies:
+    "@react-aria/focus": "npm:^3.18.3"
     "@react-aria/form": "npm:^3.0.9"
     "@react-aria/i18n": "npm:^3.12.3"
     "@react-aria/interactions": "npm:^3.22.3"