Editor: Restore focus on Escape by activeElement at focus

packages/editor/src/components/navigable-toolbar/index.js

@@ -11,22 +11,16 @@ import { Component, findDOMNode } from '@wordpress/element';
 import { focus } from '@wordpress/dom';
 import { ESCAPE } from '@wordpress/keycodes';
 
-/**
- * Browser dependencies
- */
-
-const { Node, getSelection } = window;
-
 class NavigableToolbar extends Component {
 	constructor() {
 		super( ...arguments );
 
 		this.bindNode = this.bindNode.bind( this );
 		this.focusToolbar = this.focusToolbar.bind( this );
-		this.focusSelection = this.focusSelection.bind( this );
+		this.restoreFocus = this.restoreFocus.bind( this );
 
 		this.switchOnKeyDown = cond( [
-			[ matchesProperty( [ 'keyCode' ], ESCAPE ), this.focusSelection ],
+			[ matchesProperty( [ 'keyCode' ], ESCAPE ), this.restoreFocus ],
 		] );
 	}
 
@@ -37,34 +31,37 @@ class NavigableToolbar extends Component {
 		this.toolbar = findDOMNode( ref );
 	}
 
+	/**
+	 * Shifts focus to the first tabbable element within the toolbar container,
+	 * if one exists.
+	 */
 	focusToolbar() {
 		const tabbables = focus.tabbable.find( this.toolbar );
-		if ( tabbables.length ) {
-			tabbables[ 0 ].focus();
+		if ( ! tabbables.length ) {
+			return;
 		}
+
+		// Track the original active element prior to shifting focus, so that
+		// focus can be returned if the user presses Escape while in toolbar.
+		//
+		// [TODO]: Currently, to support "stepping down" through multiple
+		// navigable toolbars, the assigned value is only cleared after
+		// pressing Escape. If the user manually shifts focus away and later
+		// returns to press Escape, it will still (wrongly) use the earlier-
+		// assigned value.
+		this.activeElementBeforeFocus = document.activeElement;
+
+		tabbables[ 0 ].focus();
 	}
 
 	/**
-	 * Programmatically shifts focus to the element where the current selection
-	 * exists, if there is a selection.
+	 * Restores focus to the active element at the time focus was
+	 * programattically shifted to the toolbar, if one exists.
 	 */
-	focusSelection() {
-		// Ensure that a selection exists.
-		const selection = getSelection();
-		if ( ! selection ) {
-			return;
-		}
-
-		// Focus node may be a text node, which cannot be focused directly.
-		// Find its parent element instead.
-		const { focusNode } = selection;
-		let focusElement = focusNode;
-		if ( focusElement.nodeType !== Node.ELEMENT_NODE ) {
-			focusElement = focusElement.parentElement;
-		}
-
-		if ( focusElement ) {
-			focusElement.focus();
+	restoreFocus() {
+		if ( this.activeElementBeforeFocus ) {
+			this.activeElementBeforeFocus.focus();
+			delete this.activeElementBeforeFocus;
 		}
 	}