Components: Add onFocusLoss option to withFocusReturn

packages/components/src/higher-order/with-focus-return/README.md

@@ -1,7 +1,13 @@
 # withFocusReturn
 
+`withFocusReturn` is a higher-order component used typically in scenarios of short-lived elements (modals, dropdowns) where, upon the element's unmounting, focus should be restored to the focused element which had initiated it being rendered.
+
+Optionally, it can be used in combination with a `FocusRenderProvider` which, when rendered toward the top of an application, will remember a history of elements focused during a session. This can provide safeguards for scenarios where one short-lived element triggers the creation of another (e.g. a dropdown menu triggering a modal display). The combined effect of `FocusRenderProvider` and `withFocusReturn` is that focus will be returned to the most recent focused element which is still present in the document.
+
 ## Usage
 
+### `withFocusReturn`
+
 ```jsx
 import { withFocusReturn, TextControl, Button } from '@wordpress/components';
 import { withState } from '@wordpress/compose';
@@ -39,3 +45,43 @@ const MyComponentWithFocusReturn = withState( {
 	);
 } ); 
 ```
+
+`withFocusReturn` can optionally be called as a higher-order function creator. Provided an options object, a new higher-order function is returned.
+
+Currently, the following options are supported:
+
+#### `onFocusReturn`
+
+An optional function which allows the developer to customize the focus return behavior. A return value of `false` should be returned from this function to indicate that the default focus return behavior should be skipped.
+
+- Type: `Function`
+- Required: No
+
+_Example:_
+
+```jsx
+function MyComponent() {
+	return <textarea />;
+}
+
+const EnhancedMyComponent = withFocusReturn( {
+	onFocusReturn() {
+		document.getElementById( 'other-input' ).focus();
+		return false;
+	},
+} )( MyComponent );
+```
+
+### `FocusReturnProvider`
+
+```jsx
+import { FocusReturnProvider } from '@wordpress/components';
+
+function App() {
+	return (
+		<FocusReturnProvider>
+			{ /* ... */ }
+		</FocusReturnProvider>
+	);
+}
+```

packages/components/src/higher-order/with-focus-return/context.js

@@ -0,0 +1,71 @@
+/**
+ * External dependencies
+ */
+import { uniq } from 'lodash';
+
+/**
+ * WordPress dependencies
+ */
+import { Component, createContext } from '@wordpress/element';
+
+const { Provider, Consumer } = createContext( {
+	focusHistory: [],
+} );
+
+Provider.displayName = 'FocusReturnProvider';
+Consumer.displayName = 'FocusReturnConsumer';
+
+/**
+ * The maximum history length to capture for the focus stack. When exceeded,
+ * items should be shifted from the stack for each consecutive push.
+ *
+ * @type {number}
+ */
+const MAX_STACK_LENGTH = 100;
+
+class FocusReturnProvider extends Component {
+	constructor() {
+		super( ...arguments );
+
+		this.onFocus = this.onFocus.bind( this );
+
+		this.state = {
+			focusHistory: [],
+		};
+	}
+
+	onFocus( event ) {
+		const { focusHistory } = this.state;
+
+		// Push the focused element to the history stack, keeping only unique
+		// members but preferring the _last_ occurrence of any duplicates.
+		// Lodash's `uniq` behavior favors the first occurrence, so the array
+		// is temporarily reversed prior to it being called upon. Uniqueness
+		// helps avoid situations where, such as in a constrained tabbing area,
+		// the user changes focus enough within a transient element that the
+		// stack may otherwise only consist of members pending destruction, at
+		// which point focus might have been lost.
+		const nextFocusHistory = uniq(
+			[ ...focusHistory, event.target ]
+				.slice( -1 * MAX_STACK_LENGTH )
+				.reverse()
+		).reverse();
+
+		this.setState( { focusHistory: nextFocusHistory } );
+	}
+
+	render() {
+		const { children, className } = this.props;
+
+		return (
+			<Provider value={ this.state }>
+				<div onFocus={ this.onFocus } className={ className }>
+					{ children }
+				</div>
+			</Provider>
+		);
+	}
+}
+
+export default FocusReturnProvider;
+export { Consumer };

packages/components/src/higher-order/with-focus-return/index.js

@@ -1,39 +1,105 @@
+/**
+ * External dependencies
+ */
+import { stubTrue, without } from 'lodash';
+
 /**
  * WordPress dependencies
  */
 import { Component } from '@wordpress/element';
 import { createHigherOrderComponent } from '@wordpress/compose';
 
+/**
+ * Internal dependencies
+ */
+import Provider, { Consumer } from './context';
+
+/**
+ * Returns true if the given object is component-like. An object is component-
+ * like if it is an instance of wp.element.Component, or is a function.
+ *
+ * @param {*} object Object to test.
+ *
+ * @return {boolean} Whether object is component-like.
+ */
+function isComponentLike( object ) {
+	return (
+		object instanceof Component ||
+		typeof object === 'function'
+	);
+}
+
 /**
  * Higher Order Component used to be used to wrap disposable elements like
  * sidebars, modals, dropdowns. When mounting the wrapped component, we track a
  * reference to the current active element so we know where to restore focus
  * when the component is unmounted.
  *
- * @param {WPElement} WrappedComponent The disposable component.
+ * @param {(WPComponent|Object)} options The component to be enhanced with
+ *                                       focus return behavior, or an object
+ *                                       describing the component and the
+ *                                       focus return characteristics.
  *
  * @return {Component} Component with the focus restauration behaviour.
  */
-export default createHigherOrderComponent(
-	( WrappedComponent ) => {
-		return class extends Component {
+function withFocusReturn( options ) {
+	// Normalize as overloaded form `withFocusReturn( options )( Component )`
+	// or as `withFocusReturn( Component )`.
+	if ( isComponentLike( options ) ) {
+		const WrappedComponent = options;
+		return withFocusReturn( {} )( WrappedComponent );
+	}
+
+	const { onFocusReturn = stubTrue } = options;
+
+	return function( WrappedComponent ) {
+		class FocusReturn extends Component {
 			constructor() {
 				super( ...arguments );
 
-				this.setIsFocusedTrue = () => this.isFocused = true;
-				this.setIsFocusedFalse = () => this.isFocused = false;
+				this.ownFocusedElements = new Set;
 				this.activeElementOnMount = document.activeElement;
+				this.setIsFocusedFalse = () => this.isFocused = false;
+				this.setIsFocusedTrue = ( event ) => {
+					this.ownFocusedElements.add( event.target );
+					this.isFocused = true;
+				};
 			}
 
 			componentWillUnmount() {
-				const { activeElementOnMount, isFocused } = this;
-				if ( ! activeElementOnMount ) {
+				const {
+					activeElementOnMount,
+					isFocused,
+					ownFocusedElements,
+				} = this;
+
+				if ( ! isFocused ) {
 					return;
 				}
 
-				const { body, activeElement } = document;
-				if ( isFocused || null === activeElement || body === activeElement ) {
-					activeElementOnMount.focus();
+				// Defer to the component's own explicit focus return behavior,
+				// if specified. The function should return `false` to prevent
+				// the default behavior otherwise occurring here. This allows
+				// for support that the `onFocusReturn` decides to allow the
+				// default behavior to occur under some conditions.
+				if ( onFocusReturn() === false ) {
+					return;
+				}
+
+				const stack = [
+					...without(
+						this.props.focusHistory,
+						...ownFocusedElements
+					),
+					activeElementOnMount,
+				];
+
+				let candidate;
+				while ( ( candidate = stack.pop() ) ) {
+					if ( document.body.contains( candidate ) ) {
+						candidate.focus();
+						return;
+					}
 				}
 			}
 
@@ -47,6 +113,15 @@ export default createHigherOrderComponent(
 					</div>
 				);
 			}
-		};
-	}, 'withFocusReturn'
-);
+		}
+
+		return ( props ) => (
+			<Consumer>
+				{ ( context ) => <FocusReturn { ...props } { ...context } /> }
+			</Consumer>
+		);
+	};
+}
+
+export default createHigherOrderComponent( withFocusReturn, 'withFocusReturn' );
+export { Provider };

packages/components/src/higher-order/with-focus-return/test/index.js

@@ -2,21 +2,22 @@
  * External dependencies
  */
 import renderer from 'react-test-renderer';
+import { mount } from 'enzyme';
 
 /**
  * WordPress dependencies
  */
-import { Component } from '@wordpress/element';
+import { Component, createElement } from '@wordpress/element';
 
 /**
  * Internal dependencies
  */
-import withFocusReturn from '../';
+import withFocusReturn, { Provider } from '../';
 
 class Test extends Component {
 	render() {
 		return (
-			<div className="test">Testing</div>
+			<div className="test"><textarea /></div>
 		);
 	}
 }
@@ -47,7 +48,7 @@ describe( 'withFocusReturn()', () => {
 			const wrappedElementShallow = wrappedElement.children[ 0 ];
 			expect( wrappedElementShallow.props.className ).toBe( 'test' );
 			expect( wrappedElementShallow.type ).toBe( 'div' );
-			expect( wrappedElementShallow.children[ 0 ] ).toBe( 'Testing' );
+			expect( wrappedElementShallow.children[ 0 ].type ).toBe( 'textarea' );
 		} );
 
 		it( 'should pass additional props through to the wrapped element', () => {
@@ -71,17 +72,47 @@ describe( 'withFocusReturn()', () => {
 			expect( document.activeElement ).toBe( switchFocusTo );
 		} );
 
-		it( 'should return focus to element associated with HOC', () => {
-			const mountedComposite = renderer.create( <Composite /> );
-			expect( getInstance( mountedComposite ).activeElementOnMount ).toBe( activeElement );
-
-			// Change activeElement.
-			document.activeElement.blur();
-			expect( document.activeElement ).toBe( document.body );
+		it( 'should switch focus back when unmounted while having focus', () => {
+			const wrapper = mount( <Composite /> );
+			wrapper.find( 'textarea' ).at( 0 ).simulate( 'focus' );
 
 			// Should return to the activeElement saved with this component.
-			mountedComposite.unmount();
+			wrapper.unmount();
 			expect( document.activeElement ).toBe( activeElement );
 		} );
+
+		it( 'should switch focus to the most recent still-available focus target', () => {
+			const container = document.createElement( 'div' );
+			document.body.appendChild( container );
+			const wrapper = mount(
+				createElement(
+					( props ) => (
+						<Provider>
+							<input name="first" />
+							{ props.renderSecondInput && <input name="second" /> }
+							{ props.renderComposite && <Composite /> }
+						</Provider>
+					),
+					{ renderSecondInput: true }
+				),
+				{ attachTo: container }
+			);
+
+			function focus( selector ) {
+				const childWrapper = wrapper.find( selector );
+				const childNode = childWrapper.getDOMNode();
+				childWrapper.simulate( 'focus', { target: childNode } );
+			}
+
+			focus( 'input[name="first"]' );
+			jest.spyOn( wrapper.find( 'input[name="first"]' ).getDOMNode(), 'focus' );
+			focus( 'input[name="second"]' );
+			wrapper.setProps( { renderComposite: true } );
+			focus( 'textarea' );
+			wrapper.setProps( { renderSecondInput: false } );
+			wrapper.setProps( { renderComposite: false } );
+
+			expect( wrapper.find( 'input[name="first"]' ).getDOMNode().focus ).toHaveBeenCalled();
+		} );
 	} );
 } );

packages/components/src/index.js

@@ -70,6 +70,6 @@ export { default as withConstrainedTabbing } from './higher-order/with-constrain
 export { default as withFallbackStyles } from './higher-order/with-fallback-styles';
 export { default as withFilters } from './higher-order/with-filters';
 export { default as withFocusOutside } from './higher-order/with-focus-outside';
-export { default as withFocusReturn } from './higher-order/with-focus-return';
+export { default as withFocusReturn, Provider as FocusReturnProvider } from './higher-order/with-focus-return';
 export { default as withNotices } from './higher-order/with-notices';
 export { default as withSpokenMessages } from './higher-order/with-spoken-messages';

packages/edit-post/src/components/layout/index.js

@@ -6,7 +6,13 @@ import classnames from 'classnames';
 /**
  * WordPress dependencies
  */
-import { Button, Popover, ScrollLock, navigateRegions } from '@wordpress/components';
+import {
+	Button,
+	Popover,
+	ScrollLock,
+	FocusReturnProvider,
+	navigateRegions,
+} from '@wordpress/components';
 import { __ } from '@wordpress/i18n';
 import { PreserveScrollInReorder } from '@wordpress/block-editor';
 import {
@@ -66,7 +72,7 @@ function Layout( {
 		tabIndex: -1,
 	};
 	return (
-		<div className={ className }>
+		<FocusReturnProvider className={ className }>
 			<FullscreenMode />
 			<BrowserURL />
 			<UnsavedChangesWarning />
@@ -126,7 +132,7 @@ function Layout( {
 			) }
 			<Popover.Slot />
 			<PluginArea />
-		</div>
+		</FocusReturnProvider>
 	);
 }