RichText: restore focus after setting selection

packages/block-editor/src/components/rich-text/toolbar-button.js

@@ -3,8 +3,17 @@
  */
 import { Fill, ToolbarButton } from '@wordpress/components';
 import { displayShortcut } from '@wordpress/keycodes';
+import { useRef } from '@wordpress/element';
+
+export function RichTextToolbarButton( {
+	name,
+	shortcutType,
+	shortcutCharacter,
+	onClick,
+	...props
+} ) {
+	const activeElement = useRef( null );
 
-export function RichTextToolbarButton( { name, shortcutType, shortcutCharacter, ...props } ) {
 	let shortcut;
 	let fillName = 'RichText.ToolbarControls';
 
@@ -21,6 +30,19 @@ export function RichTextToolbarButton( { name, shortcutType, shortcutCharacter,
 			<ToolbarButton
 				{ ...props }
 				shortcut={ shortcut }
+				extraProps={ {
+					onMouseDown: () => {
+						activeElement.current = document.activeElement;
+					},
+				} }
+				onClick={ () => {
+					onClick( ...arguments );
+
+					if ( activeElement.current ) {
+						activeElement.current.focus();
+						activeElement.current = null;
+					}
+				} }
 			/>
 		</Fill>
 	);

packages/e2e-tests/specs/__snapshots__/rich-text.test.js.snap

@@ -48,6 +48,12 @@ exports[`RichText should not lose selection direction 1`] = `
 <!-- /wp:paragraph -->"
 `;
 
+exports[`RichText should not return focus when tabbing to formatting button 1`] = `
+"<!-- wp:paragraph -->
+<p>Some <strong>bold</strong>.</p>
+<!-- /wp:paragraph -->"
+`;
+
 exports[`RichText should not undo backtick transform with backspace after selection change 1`] = `""`;
 
 exports[`RichText should not undo backtick transform with backspace after typing 1`] = `""`;
@@ -76,14 +82,14 @@ exports[`RichText should transform backtick to code 2`] = `
 <!-- /wp:paragraph -->"
 `;
 
-exports[`RichText should update internal selection after fresh focus 1`] = `
+exports[`RichText should undo backtick transform with backspace 1`] = `
 "<!-- wp:paragraph -->
-<p>1<strong>2</strong></p>
+<p>\`a\`</p>
 <!-- /wp:paragraph -->"
 `;
 
-exports[`RichText should undo backtick transform with backspace 1`] = `
+exports[`RichText should update internal selection after fresh focus 1`] = `
 "<!-- wp:paragraph -->
-<p>\`a\`</p>
+<p>1<strong>2</strong></p>
 <!-- /wp:paragraph -->"
 `;

packages/e2e-tests/specs/blocks/list.test.js

@@ -182,6 +182,7 @@ describe( 'List', () => {
 		await page.keyboard.type( 'one' );
 		await page.keyboard.press( 'Enter' );
 		await clickBlockToolbarButton( 'Indent list item' );
+		await pressKeyTimes( 'Tab', 6 );
 		await page.keyboard.type( 'two' );
 		await transformBlockTo( 'Paragraph' );
 
@@ -264,6 +265,7 @@ describe( 'List', () => {
 		await page.keyboard.type( 'one' );
 		await page.keyboard.press( 'Enter' );
 		await clickBlockToolbarButton( 'Indent list item' );
+		await pressKeyTimes( 'Tab', 6 );
 		await page.keyboard.type( 'two' );
 		await page.keyboard.press( 'Enter' );
 		await page.keyboard.type( 'three' );

packages/e2e-tests/specs/plugins/annotations.test.js

@@ -49,6 +49,7 @@ describe( 'Using Plugins API', () => {
 		// Click add annotation button.
 		const addAnnotationButton = ( await page.$x( "//button[contains(text(), 'Add annotation')]" ) )[ 0 ];
 		await addAnnotationButton.click();
+		await page.evaluate( () => document.querySelector( '[contenteditable]' ).focus() );
 	}
 
 	/**
@@ -60,6 +61,7 @@ describe( 'Using Plugins API', () => {
 		// Click remove annotations button.
 		const addAnnotationButton = ( await page.$x( "//button[contains(text(), 'Remove annotations')]" ) )[ 0 ];
 		await addAnnotationButton.click();
+		await page.evaluate( () => document.querySelector( '[contenteditable]' ).focus() );
 	}
 
 	/**

packages/e2e-tests/specs/rich-text.test.js

@@ -7,6 +7,7 @@ import {
 	insertBlock,
 	clickBlockAppender,
 	pressKeyWithModifier,
+	pressKeyTimes,
 } from '@wordpress/e2e-test-utils';
 
 describe( 'RichText', () => {
@@ -92,6 +93,30 @@ describe( 'RichText', () => {
 		expect( await getEditedPostContent() ).toMatchSnapshot();
 	} );
 
+	it( 'should not return focus when tabbing to formatting button', async () => {
+		await clickBlockAppender();
+		await page.keyboard.type( 'Some ' );
+		await pressKeyWithModifier( 'alt', 'F10' );
+		// Wait for button to receive focus.
+		await page.waitForFunction( () =>
+			document.activeElement.nodeName === 'BUTTON'
+		);
+		await pressKeyTimes( 'Tab', 2 );
+		await page.keyboard.press( 'Space' );
+		await pressKeyTimes( 'Tab', 5 );
+		await page.keyboard.type( 'bold' );
+		await pressKeyWithModifier( 'alt', 'F10' );
+		await page.waitForFunction( () =>
+			document.activeElement.nodeName === 'BUTTON'
+		);
+		await pressKeyTimes( 'Tab', 2 );
+		await page.keyboard.press( 'Space' );
+		await pressKeyTimes( 'Tab', 5 );
+		await page.keyboard.type( '.' );
+
+		expect( await getEditedPostContent() ).toMatchSnapshot();
+	} );
+
 	it( 'should transform backtick to code', async () => {
 		await clickBlockAppender();
 		await page.keyboard.type( 'A `backtick`' );

packages/format-library/src/link/inline.js

@@ -149,9 +149,9 @@ class InlineLinkUI extends Component {
 
 		if ( isCollapsed( value ) && ! isActive ) {
 			const toInsert = applyFormat( create( { text: url } ), format, 0, url.length );
-			onChange( insert( value, toInsert ) );
+			onChange( insert( value, toInsert ), { focus: true } );
 		} else {
-			onChange( applyFormat( value, format ) );
+			onChange( applyFormat( value, format ), { focus: true } );
 		}
 
 		this.resetState();

packages/rich-text/src/component/index.js

@@ -170,6 +170,8 @@ class RichText extends Component {
 	applyRecord( record, { domOnly } = {} ) {
 		const { __unstableMultilineTag: multilineTag } = this.props;
 
+		this.ignoreFocus = true;
+
 		apply( {
 			value: record,
 			current: this.editableRef,
@@ -179,6 +181,8 @@ class RichText extends Component {
 			__unstableDomOnly: domOnly,
 			placeholder: this.props.placeholder,
 		} );
+
+		this.ignoreFocus = false;
 	}
 
 	/**
@@ -276,25 +280,31 @@ class RichText extends Component {
 	onFocus() {
 		const { unstableOnFocus } = this.props;
 
+		if ( this.ignoreFocus ) {
+			return;
+		}
+
 		if ( unstableOnFocus ) {
 			unstableOnFocus();
 		}
 
 		this.recalculateBoundaryStyle();
 
-		// We know for certain that on focus, the old selection is invalid. It
-		// will be recalculated on the next mouseup, keyup, or touchend event.
-		const index = undefined;
-		const activeFormats = undefined;
+		if ( ! this.props.__unstableIsSelected ) {
+			// We know for certain that on focus, the old selection is invalid. It
+			// will be recalculated on the next mouseup, keyup, or touchend event.
+			const index = undefined;
+			const activeFormats = undefined;
 
-		this.record = {
-			...this.record,
-			start: index,
-			end: index,
-			activeFormats,
-		};
-		this.props.onSelectionChange( index, index );
-		this.setState( { activeFormats } );
+			this.record = {
+				...this.record,
+				start: index,
+				end: index,
+				activeFormats,
+			};
+			this.props.onSelectionChange( index, index );
+			this.setState( { activeFormats } );
+		}
 
 		// Update selection as soon as possible, which is at the next animation
 		// frame. The event listener for selection changes may be added too late
@@ -524,8 +534,14 @@ class RichText extends Component {
 	 * @param {Object}  $2                Named options.
 	 * @param {boolean} $2.withoutHistory If true, no undo level will be
 	 *                                    created.
+	 * @param {boolean} $2.focus          If true, the rich text element will
+	 *                                    receive focus.
 	 */
-	onChange( record, { withoutHistory } = {} ) {
+	onChange( record, { withoutHistory, focus } = {} ) {
+		if ( focus ) {
+			this.editableRef.focus();
+		}
+
 		this.applyRecord( record );
 
 		const { start, end, activeFormats = [] } = record;

packages/rich-text/src/to-dom.js

@@ -278,10 +278,7 @@ export function applySelection( { startPath, endPath }, current ) {
 	range.setStart( startContainer, startOffset );
 	range.setEnd( endContainer, endOffset );
 
-	// Set back focus if focus is lost.
-	if ( ownerDocument.activeElement !== current ) {
-		current.focus();
-	}
+	const { activeElement } = ownerDocument;
 
 	if ( selection.rangeCount > 0 ) {
 		// If the to be added range and the live range are the same, there's no
@@ -294,4 +291,5 @@ export function applySelection( { startPath, endPath }, current ) {
 	}
 
 	selection.addRange( range );
+	activeElement.focus();
 }