add use tree walker hook (#316)

* add useTreeWalker hooks

We got a PR to fix the createTreeWalker so that it also works in IE11.
We don't actively support IE11, so if things work (with polyfills) then
it's good but I don't want to maintain IE11 specific code.

That said, I wanted to abstract the createTreeWalker code to a nice
little hook. The fix for IE is also pretty small, it uses a function
instead of an object and it has a last argument that is deprecated, but
has no obvious effect for our use cases.

Since the incoming PR was based on the `main` branch (where we only had
1 reference to createTreeWalker), I wanted to make sure that we got all
the references on the latest `develop` branch.

Closes: #295
Co-authored-by: Simon VDB <[email protected]>

* use useTreeWalker hook

Co-authored-by: Simon VDB <[email protected]>

Author: RobinMalfait

packages/@headlessui-react/src/components/menu/menu.tsx

@@ -32,6 +32,7 @@ import { Focus, calculateActiveIndex } from '../../utils/calculate-active-index'
 import { isDisabledReactIssue7711 } from '../../utils/bugs'
 import { isFocusableElement, FocusableMode } from '../../utils/focus-management'
 import { useWindowEvent } from '../../hooks/use-window-event'
+import { useTreeWalker } from '../../hooks/use-tree-walker'
 
 enum MenuStates {
   Open,
@@ -338,23 +339,18 @@ let Items = forwardRefWithAs(function Items<TTag extends ElementType = typeof DE
     container.focus({ preventScroll: true })
   }, [state.menuState, state.itemsRef])
 
-  useIsoMorphicEffect(() => {
-    let container = state.itemsRef.current
-    if (!container) return
-    if (state.menuState !== MenuStates.Open) return
-
-    let walker = document.createTreeWalker(container, NodeFilter.SHOW_ELEMENT, {
-      acceptNode(node: HTMLElement) {
-        if (node.getAttribute('role') === 'menuitem') return NodeFilter.FILTER_REJECT
-        if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
-        return NodeFilter.FILTER_ACCEPT
-      },
-    })
-
-    while (walker.nextNode()) {
-      ;(walker.currentNode as HTMLElement).setAttribute('role', 'none')
-    }
-  }, [state.menuState, state.itemsRef])
+  useTreeWalker({
+    container: state.itemsRef.current,
+    enabled: state.menuState === MenuStates.Open,
+    accept(node) {
+      if (node.getAttribute('role') === 'menuitem') return NodeFilter.FILTER_REJECT
+      if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
+      return NodeFilter.FILTER_ACCEPT
+    },
+    walk(node) {
+      node.setAttribute('role', 'none')
+    },
+  })
 
   let handleKeyDown = useCallback(
     (event: ReactKeyboardEvent<HTMLDivElement>) => {

packages/@headlessui-react/src/components/radio-group/radio-group.tsx

@@ -23,6 +23,7 @@ import { focusIn, Focus, FocusResult } from '../../utils/focus-management'
 import { useFlags } from '../../hooks/use-flags'
 import { Label, useLabels } from '../../components/label/label'
 import { Description, useDescriptions } from '../../components/description/description'
+import { useTreeWalker } from '../../hooks/use-tree-walker'
 
 interface Option {
   id: string
@@ -131,22 +132,17 @@ export function RadioGroup<
     [onChange, value]
   )
 
-  useIsoMorphicEffect(() => {
-    let container = radioGroupRef.current
-    if (!container) return
-
-    let walker = document.createTreeWalker(container, NodeFilter.SHOW_ELEMENT, {
-      acceptNode(node: HTMLElement) {
-        if (node.getAttribute('role') === 'radio') return NodeFilter.FILTER_REJECT
-        if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
-        return NodeFilter.FILTER_ACCEPT
-      },
-    })
-
-    while (walker.nextNode()) {
-      ;(walker.currentNode as HTMLElement).setAttribute('role', 'none')
-    }
-  }, [radioGroupRef])
+  useTreeWalker({
+    container: radioGroupRef.current,
+    accept(node) {
+      if (node.getAttribute('role') === 'radio') return NodeFilter.FILTER_REJECT
+      if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
+      return NodeFilter.FILTER_ACCEPT
+    },
+    walk(node) {
+      node.setAttribute('role', 'none')
+    },
+  })
 
   let handleKeyDown = useCallback(
     (event: ReactKeyboardEvent<HTMLButtonElement>) => {

packages/@headlessui-react/src/hooks/use-tree-walker.ts

@@ -0,0 +1,42 @@
+import { useRef, useEffect } from 'react'
+import { useIsoMorphicEffect } from './use-iso-morphic-effect'
+
+type AcceptNode = (
+  node: HTMLElement
+) =>
+  | typeof NodeFilter.FILTER_ACCEPT
+  | typeof NodeFilter.FILTER_SKIP
+  | typeof NodeFilter.FILTER_REJECT
+
+export function useTreeWalker({
+  container,
+  accept,
+  walk,
+  enabled = true,
+}: {
+  container: HTMLElement | null
+  accept: AcceptNode
+  walk(node: HTMLElement): void
+  enabled?: boolean
+}) {
+  let acceptRef = useRef(accept)
+  let walkRef = useRef(walk)
+
+  useEffect(() => {
+    acceptRef.current = accept
+    walkRef.current = walk
+  }, [accept, walk])
+
+  useIsoMorphicEffect(() => {
+    if (!container) return
+    if (!enabled) return
+
+    let accept = acceptRef.current
+    let walk = walkRef.current
+
+    let acceptNode = Object.assign((node: HTMLElement) => accept(node), { acceptNode: accept })
+    let walker = document.createTreeWalker(container, NodeFilter.SHOW_ELEMENT, acceptNode, false)
+
+    while (walker.nextNode()) walk(walker.currentNode as HTMLElement)
+  }, [container, enabled, acceptRef, walkRef])
+}

packages/@headlessui-vue/src/components/menu/menu.ts

@@ -18,6 +18,7 @@ import { Focus, calculateActiveIndex } from '../../utils/calculate-active-index'
 import { resolvePropValue } from '../../utils/resolve-prop-value'
 import { dom } from '../../utils/dom'
 import { useWindowEvent } from '../../hooks/use-window-event'
+import { useTreeWalker } from '../../hooks/use-tree-walker'
 
 enum MenuStates {
   Open,
@@ -291,22 +292,17 @@ export let MenuItems = defineComponent({
     let id = `headlessui-menu-items-${useId()}`
     let searchDebounce = ref<ReturnType<typeof setTimeout> | null>(null)
 
-    watchEffect(() => {
-      let container = dom(api.itemsRef)
-      if (!container) return
-      if (api.menuState.value !== MenuStates.Open) return
-
-      let walker = document.createTreeWalker(container, NodeFilter.SHOW_ELEMENT, {
-        acceptNode(node: HTMLElement) {
-          if (node.getAttribute('role') === 'menuitem') return NodeFilter.FILTER_REJECT
-          if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
-          return NodeFilter.FILTER_ACCEPT
-        },
-      })
-
-      while (walker.nextNode()) {
-        ;(walker.currentNode as HTMLElement).setAttribute('role', 'none')
-      }
+    useTreeWalker({
+      container: computed(() => dom(api.itemsRef)),
+      enabled: computed(() => api.menuState.value === MenuStates.Open),
+      accept(node) {
+        if (node.getAttribute('role') === 'menuitem') return NodeFilter.FILTER_REJECT
+        if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
+        return NodeFilter.FILTER_ACCEPT
+      },
+      walk(node) {
+        node.setAttribute('role', 'none')
+      },
     })
 
     function handleKeyDown(event: KeyboardEvent) {

packages/@headlessui-vue/src/components/radio-group/radio-group.ts

@@ -7,7 +7,6 @@ import {
   provide,
   ref,
   toRaw,
-  watchEffect,
 
   // Types
   InjectionKey,
@@ -22,6 +21,7 @@ import { render } from '../../utils/render'
 import { Label, useLabels } from '../label/label'
 import { Description, useDescriptions } from '../description/description'
 import { resolvePropValue } from '../../utils/resolve-prop-value'
+import { useTreeWalker } from '../../hooks/use-tree-walker'
 
 interface Option {
   id: string
@@ -121,21 +121,16 @@ export let RadioGroup = defineComponent({
     // @ts-expect-error ...
     provide(RadioGroupContext, api)
 
-    watchEffect(() => {
-      let container = dom(radioGroupRef)
-      if (!container) return
-
-      let walker = document.createTreeWalker(container, NodeFilter.SHOW_ELEMENT, {
-        acceptNode(node: HTMLElement) {
-          if (node.getAttribute('role') === 'radio') return NodeFilter.FILTER_REJECT
-          if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
-          return NodeFilter.FILTER_ACCEPT
-        },
-      })
-
-      while (walker.nextNode()) {
-        ;(walker.currentNode as HTMLElement).setAttribute('role', 'none')
-      }
+    useTreeWalker({
+      container: computed(() => dom(radioGroupRef)),
+      accept(node) {
+        if (node.getAttribute('role') === 'radio') return NodeFilter.FILTER_REJECT
+        if (node.hasAttribute('role')) return NodeFilter.FILTER_SKIP
+        return NodeFilter.FILTER_ACCEPT
+      },
+      walk(node) {
+        node.setAttribute('role', 'none')
+      },
     })
 
     function handleKeyDown(event: KeyboardEvent) {

packages/@headlessui-vue/src/hooks/use-tree-walker.ts

@@ -0,0 +1,31 @@
+import { watchEffect, ComputedRef } from 'vue'
+
+type AcceptNode = (
+  node: HTMLElement
+) =>
+  | typeof NodeFilter.FILTER_ACCEPT
+  | typeof NodeFilter.FILTER_SKIP
+  | typeof NodeFilter.FILTER_REJECT
+
+export function useTreeWalker({
+  container,
+  accept,
+  walk,
+  enabled,
+}: {
+  container: ComputedRef<HTMLElement | null>
+  accept: AcceptNode
+  walk(node: HTMLElement): void
+  enabled?: ComputedRef<boolean>
+}) {
+  watchEffect(() => {
+    let root = container.value
+    if (!root) return
+    if (enabled !== undefined && !enabled.value) return
+
+    let acceptNode = Object.assign((node: HTMLElement) => accept(node), { acceptNode: accept })
+    let walker = document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT, acceptNode, false)
+
+    while (walker.nextNode()) walk(walker.currentNode as HTMLElement)
+  })
+}