Refactor setting focusedBeforeMount ref

jeryj · ajlende · web-flow · commit 0c834eb68038 · 2024-12-18T11:53:46.000-06:00
Co-authored-by: Alex Lende &lt;alex+github.com@lende.xyz&gt;
diff --git a/packages/compose/src/hooks/use-focus-return/index.js b/packages/compose/src/hooks/use-focus-return/index.js
@@ -48,21 +48,13 @@ function useFocusReturn( onFocusReturn ) {
 				return;
 			}
 
-			/** @type {any} */
-			const activeElement = node.ownerDocument.activeElement;
+			const activeDocument =
+				node.ownerDocument.activeElement instanceof
+				window.HTMLIFrameElement
+					? node.ownerDocument.activeElement.contentDocument
+					: node.ownerDocument;
 
-			// If the activeElement is an iframe, we need to get the active element within the iframe.
-			// Otherwise, focus from items within the iframed canvas will get sent to the iframe itself,
-			// not the active element within the iframe.
-			if (
-				activeElement?.tagName === 'IFRAME' &&
-				activeElement.contentDocument?.activeElement
-			) {
-				focusedBeforeMount.current =
-					activeElement.contentDocument.activeElement;
-			} else {
-				focusedBeforeMount.current = activeElement;
-			}
+			focusedBeforeMount.current = activeDocument?.activeElement ?? null;
 		} else if ( focusedBeforeMount.current ) {
 			const isFocused = ref.current?.contains(
 				ref.current?.ownerDocument.activeElement
