From 7d5fa9db48470c61f3058644918ad00bc97819a8 Mon Sep 17 00:00:00 2001
From: sarayourfriend <24264157+sarayourfriend@users.noreply.github.com>
Date: Fri, 21 May 2021 08:09:11 -0700
Subject: [PATCH] compose: Add types to `useFocusReturn` (#31949)

* compose: Add types to `useFocusReturn`

* Slightly improve types

* Use null as default value

* Fix missing null type.

Co-authored-by: Marco Ciampini <marco.ciampo@gmail.com>

Co-authored-by: Marco Ciampini <marco.ciampo@gmail.com>
---
 packages/compose/README.md                     |  4 ++--
 .../src/hooks/use-focus-return/index.js        | 18 ++++++++++--------
 packages/compose/tsconfig.json                 |  1 +
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/packages/compose/README.md b/packages/compose/README.md
index 3f92504fdc3f7f..f72dfb4ce9e810 100644
--- a/packages/compose/README.md
+++ b/packages/compose/README.md
@@ -261,11 +261,11 @@ const WithFocusReturn = () => {
 
 _Parameters_
 
--   _onFocusReturn_ `Function?`: Overrides the default return behavior.
+-   _onFocusReturn_ `[() => void]`: Overrides the default return behavior.
 
 _Returns_
 
--   `Function`: Element Ref.
+-   `import('react').RefCallback<HTMLElement>`: Element Ref.
 
 <a name="useInstanceId" href="#useInstanceId">#</a> **useInstanceId**
 
diff --git a/packages/compose/src/hooks/use-focus-return/index.js b/packages/compose/src/hooks/use-focus-return/index.js
index d7c5975d0ee192..5f16a76167648f 100644
--- a/packages/compose/src/hooks/use-focus-return/index.js
+++ b/packages/compose/src/hooks/use-focus-return/index.js
@@ -9,8 +9,8 @@ import { useRef, useEffect, useCallback } from '@wordpress/element';
  * previously focused element when closed.
  * The current hook implements the returning behavior.
  *
- * @param {Function?} onFocusReturn Overrides the default return behavior.
- * @return {Function} Element Ref.
+ * @param {() => void} [onFocusReturn] Overrides the default return behavior.
+ * @return {import('react').RefCallback<HTMLElement>} Element Ref.
  *
  * @example
  * ```js
@@ -28,8 +28,10 @@ import { useRef, useEffect, useCallback } from '@wordpress/element';
  * ```
  */
 function useFocusReturn( onFocusReturn ) {
-	const ref = useRef();
-	const focusedBeforeMount = useRef();
+	/** @type {import('react').MutableRefObject<null | HTMLElement>} */
+	const ref = useRef( null );
+	/** @type {import('react').MutableRefObject<null | Element>} */
+	const focusedBeforeMount = useRef( null );
 	const onFocusReturnRef = useRef( onFocusReturn );
 	useEffect( () => {
 		onFocusReturnRef.current = onFocusReturn;
@@ -47,11 +49,11 @@ function useFocusReturn( onFocusReturn ) {
 
 			focusedBeforeMount.current = node.ownerDocument.activeElement;
 		} else if ( focusedBeforeMount.current ) {
-			const isFocused = ref.current.contains(
-				ref.current.ownerDocument.activeElement
+			const isFocused = ref.current?.contains(
+				ref.current?.ownerDocument.activeElement
 			);
 
-			if ( ref.current.isConnected && ! isFocused ) {
+			if ( ref.current?.isConnected && ! isFocused ) {
 				return;
 			}
 
@@ -62,7 +64,7 @@ function useFocusReturn( onFocusReturn ) {
 			if ( onFocusReturnRef.current ) {
 				onFocusReturnRef.current();
 			} else {
-				focusedBeforeMount.current.focus();
+				/** @type {null | HTMLElement} */ ( focusedBeforeMount.current )?.focus();
 			}
 		}
 	}, [] );
diff --git a/packages/compose/tsconfig.json b/packages/compose/tsconfig.json
index 6bc5f43325c586..9544096d8ff048 100644
--- a/packages/compose/tsconfig.json
+++ b/packages/compose/tsconfig.json
@@ -18,6 +18,7 @@
 		"src/hooks/use-constrained-tabbing/**/*",
 		"src/hooks/use-debounce/**/*",
 		"src/hooks/use-instance-id/**/*",
+		"src/hooks/use-focus-return/**/*",
 		"src/utils/**/*"
 	]
 }