feat: smart return focus feature

README.md

@@ -61,8 +61,9 @@ Demo - https://codesandbox.io/s/5wmrwlvxv4.
  FocusLock has few props to tune behavior, all props are optional:
   - `disabled`, to disable(enable) behavior without altering the tree.
   - `className`, to set the `className` of the internal wrapper.
-  - `returnFocus`, to return focus into initial position on unmount(not disable).
-> By default `returnFocus` is disabled, so FocusLock will __not__ restore original focus on deactivation.
+  - `returnFocus`, to return focus into initial position on unmount
+> By default `returnFocus` is disabled, so FocusLock __will not__ restore original focus on deactivation.
+> This was done mostly to avoid breaking changes. We __strong recommend enabling it__, to provide a better user experience.
     
   This is expected behavior for Modals, but it is better to implement it by your self. See [unmounting and focus management](https://github.com/theKashey/react-focus-lock#unmounting-and-focus-management) for details
   - `persistentFocus=false`, requires any element to be focused. This also disables text selections inside, and __outside__ focus lock.
@@ -329,6 +330,32 @@ to allow user _tab_ into address bar.
 >
 ```
 
+## Return focus to another node
+In some cases the original node that was focused before the lock was activated is not the desired node to return focus to.
+Some times this node might not exists at all.
+
+- first of all, FocusLock need a moment to record this node, please do not hide it onClick, but hide onBlur (Dropdown, looking at you)
+- second, you may specify a callback as `returnFocus`, letting you decide where to return focus to.
+```tsx
+<FocusLock
+    returnFocus={(suggestedNode) => {
+        // somehow activeElement should not be changed
+        if(document.activeElement.hasAttributes('main-content')) {
+            // opt out from default behavior
+            return false;
+        }
+        if (someCondition(suggestedNode)) {
+            // proceed with the suggested node
+            return true;
+        } 
+        // handle return focus manually
+        document.getElementById('the-button').focus();
+        // opt out from default behavior
+        return false;
+    }}
+/>
+````
+
 ## Return focus with no scroll
 > read more at the [issue #83](https://github.com/theKashey/react-focus-lock/issues/83) or
 [mdn article](https://developer.mozilla.org/en-US/docs/Web/API/HTMLElement/focus).

_tests/FocusLock.spec.js

@@ -244,6 +244,81 @@ describe('react-focus-lock', () => {
       expect(document.activeElement.innerHTML).to.be.equal('d-action0');
     });
 
+    it('Should return focus to the possible place', async () => {
+      const LockTest = ({ action }) => (
+        <FocusLock returnFocus>
+          <button id="focus-action" onClick={action}>
+            inside
+          </button>
+        </FocusLock>
+      );
+
+      const TriggerTest = () => {
+        const [clicked, setClicked] = React.useState(false);
+        const [removed, setRemoved] = React.useState(false);
+        return (
+          <>
+            {removed ? null : <button id="trigger" onClick={() => setClicked(true)}>trigger</button>}
+            <button id="follower">another action</button>
+            {clicked && (
+            <LockTest
+              action={() => {
+                setRemoved(true);
+                setTimeout(() => {
+                  setClicked(false);
+                }, 1);
+              }}
+            />
+            )}
+          </>
+        );
+      };
+
+      const wrapper = mount(<TriggerTest />);
+
+      document.getElementById('trigger').focus();
+      wrapper.find('#trigger').simulate('click');
+      expect(document.activeElement.innerHTML).to.be.equal('inside');
+      // await tick(1);
+      wrapper.find('#focus-action').simulate('click');
+      await tick(5);
+      expect(document.activeElement.innerHTML).to.be.equal('another action');
+    });
+
+    it.only('Should return focus to the possible place: timing', async () => {
+      const LockTest = ({ action }) => (
+        <FocusLock returnFocus>
+          <button id="focus-action" onClick={action}>
+            inside
+          </button>
+        </FocusLock>
+      );
+
+      const TriggerTest = () => {
+        const [clicked, setClicked] = React.useState(false);
+        return (
+          <>
+            {clicked ? null : <button id="trigger" onClick={() => setClicked(true)}>trigger</button>}
+            <button id="follower">another action</button>
+            {clicked && (
+              <LockTest
+                action={() => { setClicked(false); }}
+              />
+            )}
+          </>
+        );
+      };
+
+      const wrapper = mount(<TriggerTest />);
+
+      wrapper.find('#trigger').simulate('click');
+      await tick();
+      expect(document.activeElement.innerHTML).to.be.equal('inside');
+      wrapper.find('#focus-action').simulate('click');
+      await tick();
+      expect(document.activeElement).to.be.equal(document.body);
+    });
+
     it('Should focus on inputs', (done) => {
       const wrapper = mount(<div>
         <div>

_tests/restore-focus.sidecar.spec.js

@@ -0,0 +1,51 @@
+import * as React from 'react';
+import { expect } from 'chai';
+import { render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import { sidecar } from 'use-sidecar';
+import FocusLock from '../src/UI';
+
+const tick = (tm = 1) => new Promise(resolve => setTimeout(resolve, tm));
+
+const FocusLockSidecar = sidecar(
+  () => import('../src/sidecar').then(async (x) => {
+    await tick();
+    return x;
+  }),
+);
+
+it('Should return focus to the possible place: timing', async () => {
+  const LockTest = ({ action }) => (
+    <FocusLock returnFocus sideCar={FocusLockSidecar}>
+      <button data-testid="focus-action" onClick={action}>
+        inside
+      </button>
+    </FocusLock>
+  );
+
+  const TriggerTest = () => {
+    const [clicked, setClicked] = React.useState(false);
+    return (
+      <>
+        {clicked ? null : <button data-testid="trigger" onClick={() => setClicked(true)}>trigger</button>}
+        <button id="follower">another action</button>
+        {clicked && (
+        <LockTest
+          action={() => { setClicked(false); }}
+        />
+        )}
+      </>
+    );
+  };
+
+  render(<TriggerTest />);
+
+  screen.getByTestId('trigger').focus();
+  await userEvent.click(screen.getByTestId('trigger'));
+  await tick(5);
+  expect(document.activeElement.innerHTML).to.be.equal('inside');
+  await userEvent.click(screen.getByTestId('focus-action'));
+  await tick();
+  console.log('active is ', document.activeElement.tagName);
+  expect(document.activeElement).to.be.equal(document.body);
+});

package.json

@@ -71,6 +71,7 @@
     "@storybook/addon-links": "^5.1.8",
     "@storybook/react": "^5.1.8",
     "@testing-library/react": "^12.0.0",
+    "@testing-library/user-event": "^12.0.0",
     "@types/react": "^18.0.8",
     "babel-eslint": "^10.0.1",
     "babel-loader": "^8.0.4",

src/Lock.js

@@ -7,7 +7,9 @@ import { useMergeRefs } from 'use-callback-ref';
 
 import { useEffect } from 'react';
 import { hiddenGuard } from './FocusGuard';
-import { mediumFocus, mediumBlur, mediumSidecar } from './medium';
+import {
+  mediumFocus, mediumBlur, mediumSidecar,
+} from './medium';
 import { focusScope } from './scope';
 
 const emptyArray = [];
@@ -48,10 +50,20 @@ const FocusLock = React.forwardRef(function FocusLockUI(props, parentRef) {
 
   // SIDE EFFECT CALLBACKS
 
-  const onActivation = React.useCallback(() => {
-    originalFocusedElement.current = (
-      originalFocusedElement.current || (document && document.activeElement)
-    );
+  const onActivation = React.useCallback(({ captureFocusRestore }) => {
+    if (!originalFocusedElement.current) {
+      const activeElement = document?.activeElement;
+      originalFocusedElement.current = activeElement;
+      // store stack reference
+      if (activeElement !== document.body) {
+        if (document.contains(activeElement)) {
+          originalFocusedElement.current = captureFocusRestore(activeElement);
+        } else if (shouldReturnFocus) {
+          console.error('FocusLock: returnFocus element has been removed from the DOM before stack capture. Element:', activeElement);
+        }
+      }
+    }
+
     if (observed.current && onActivationCallback) {
       onActivationCallback(observed.current);
     }
@@ -67,17 +79,10 @@ const FocusLock = React.forwardRef(function FocusLockUI(props, parentRef) {
     update();
   }, [onDeactivationCallback]);
 
-  useEffect(() => {
-    if (!disabled) {
-      // cleanup return focus on trap deactivation
-      // sideEffect/returnFocus should happen by this time
-      originalFocusedElement.current = null;
-    }
-  }, []);
-
   const returnFocus = React.useCallback((allowDefer) => {
-    const { current: returnFocusTo } = originalFocusedElement;
-    if (returnFocusTo && returnFocusTo.focus) {
+    const { current: focusRestore } = originalFocusedElement;
+    if (focusRestore) {
+      const returnFocusTo = (typeof focusRestore === 'function' ? focusRestore() : focusRestore) || document.body;
       const howToReturnFocus = typeof shouldReturnFocus === 'function' ? shouldReturnFocus(returnFocusTo) : shouldReturnFocus;
       if (howToReturnFocus) {
         const returnFocusOptions = typeof howToReturnFocus === 'object' ? howToReturnFocus : undefined;

src/Trap.js

@@ -7,6 +7,7 @@ import {
   focusIsHidden, expandFocusableNodes,
   focusNextElement,
   focusPrevElement,
+  captureFocusRestore,
 } from 'focus-lock';
 import { deferAction, extractRef } from './util';
 import { mediumFocus, mediumBlur, mediumEffect } from './medium';
@@ -211,6 +212,14 @@ function reducePropsToState(propsList) {
     .filter(({ disabled }) => !disabled);
 }
 
+const focusLockAPI = {
+  moveFocusInside,
+  focusInside,
+  focusNextElement,
+  focusPrevElement,
+  captureFocusRestore,
+};
+
 function handleStateChangeOnClient(traps) {
   const trap = traps.slice(-1)[0];
   if (trap && !lastActiveTrap) {
@@ -234,7 +243,7 @@ function handleStateChangeOnClient(traps) {
   if (trap) {
     lastActiveFocus = null;
     if (!sameTrap || lastTrap.observed !== trap.observed) {
-      trap.onActivation();
+      trap.onActivation(focusLockAPI);
     }
     activateTrap(true);
     deferAction(activateTrap);
@@ -247,12 +256,7 @@ function handleStateChangeOnClient(traps) {
 // bind medium
 mediumFocus.assignSyncMedium(onFocus);
 mediumBlur.assignMedium(onBlur);
-mediumEffect.assignMedium(cb => cb({
-  moveFocusInside,
-  focusInside,
-  focusNextElement,
-  focusPrevElement,
-}));
+mediumEffect.assignMedium(cb => cb(focusLockAPI));
 
 export default withSideEffect(
   reducePropsToState,

yarn.lock

@@ -1856,6 +1856,13 @@
     "@testing-library/dom" "^8.0.0"
     "@types/react-dom" "<18.0.0"
 
+"@testing-library/user-event@^12.0.0":
+  version "12.8.3"
+  resolved "https://registry.yarnpkg.com/@testing-library/user-event/-/user-event-12.8.3.tgz#1aa3ed4b9f79340a1e1836bc7f57c501e838704a"
+  integrity sha512-IR0iWbFkgd56Bu5ZI/ej8yQwrkCv8Qydx6RzwbKz9faXazR/+5tvYKsZQgyXJiwgpcva127YO6JcWy7YlCfofQ==
+  dependencies:
+    "@babel/runtime" "^7.12.5"
+
 "@tootallnate/once@1":
   version "1.1.2"
   resolved "https://registry.yarnpkg.com/@tootallnate/once/-/once-1.1.2.tgz#ccb91445360179a04e7fe6aff78c00ffc1eeaf82"