feat(treeview): add checked state support

Author: segunadebayo

.changeset/khaki-times-pick.md

@@ -0,0 +1,7 @@
+---
+"@zag-js/tree-view": minor
+---
+
+- Add support for checkbox state for checkbox trees via `defaultCheckedValue`, `checkedValue`, `onCheckedChange` props
+- Add callback for when `loadChildren` fails via `onLoadChildrenError` prop
+- Add `api.getCheckedMap` method to get the checked state of all nodes

.changeset/ninety-needles-joke.md

@@ -0,0 +1,5 @@
+---
+"@zag-js/collection": patch
+---
+
+Add support for `getDescendantNodes` and `getDescendantValues`

examples/next-ts/pages/compositions

@@ -5,18 +5,23 @@ export { machine } from "./tree-view.machine"
 export * from "./tree-view.props"
 export type {
   TreeViewApi as Api,
-  TreeViewProps as Props,
+  CheckedChangeDetails,
+  CheckedState,
+  CheckedValueMap,
   ElementIds,
   ExpandedChangeDetails,
   FocusChangeDetails,
+  LoadChildrenCompleteDetails,
+  LoadChildrenDetails,
+  LoadChildrenErrorDetails,
+  TreeViewMachine as Machine,
   NodeProps,
   NodeState,
+  NodeWithError,
+  TreeViewProps as Props,
   SelectionChangeDetails,
   TreeViewService as Service,
-  TreeViewMachine as Machine,
-  TreeNode,
-  LoadChildrenDetails,
-  LoadChildrenCompleteDetails,
   TreeLoadingStatus,
   TreeLoadingStatusMap,
+  TreeNode,
 } from "./tree-view.types"

packages/machines/tree-view/src/index.ts

@@ -7,6 +7,7 @@ export const anatomy = createAnatomy("tree-view").parts(
   "item",
   "itemIndicator",
   "itemText",
+  "itemCheckbox",
   "branch",
   "branchControl",
   "branchTrigger",

packages/machines/tree-view/src/tree-view.anatomy.ts

@@ -10,13 +10,20 @@ import type { EventKeyMap, NormalizeProps, PropTypes } from "@zag-js/types"
 import { add, isEqual, uniq } from "@zag-js/utils"
 import { parts } from "./tree-view.anatomy"
 import * as dom from "./tree-view.dom"
-import type { NodeProps, NodeState, TreeViewApi, TreeViewService } from "./tree-view.types"
+import type { NodeProps, NodeState, TreeNode, TreeViewApi, TreeViewService } from "./tree-view.types"
+import { getCheckedState, getCheckedValueMap } from "./utils/checked-state"
 
-export function connect<T extends PropTypes>(service: TreeViewService, normalize: NormalizeProps<T>): TreeViewApi<T> {
+export function connect<T extends PropTypes, V extends TreeNode = TreeNode>(
+  service: TreeViewService<V>,
+  normalize: NormalizeProps<T>,
+): TreeViewApi<T, V> {
   const { context, scope, computed, prop, send } = service
   const collection = prop("collection")
+
   const expandedValue = Array.from(context.get("expandedValue"))
   const selectedValue = Array.from(context.get("selectedValue"))
+  const checkedValue = Array.from(context.get("checkedValue"))
+
   const isTypingAhead = computed("isTypingAhead")
   const focusedValue = context.get("focusedValue")
   const loadingStatus = context.get("loadingStatus")
@@ -35,13 +42,29 @@ export function connect<T extends PropTypes>(service: TreeViewService, normalize
       loading: loadingStatus[value] === "loading",
       depth: indexPath.length,
       isBranch: collection.isBranchNode(node),
+      get checked() {
+        return getCheckedState(collection, node, checkedValue)
+      },
     }
   }
 
   return {
     collection,
     expandedValue,
     selectedValue,
+    checkedValue,
+    toggleChecked(value, isBranch) {
+      send({ type: "CHECKED.TOGGLE", value, isBranch })
+    },
+    setChecked(value) {
+      send({ type: "CHECKED.SET", value })
+    },
+    clearChecked() {
+      send({ type: "CHECKED.CLEAR" })
+    },
+    getCheckedMap() {
+      return getCheckedValueMap(collection, checkedValue)
+    },
     expand(value) {
       send({ type: value ? "BRANCH.EXPAND" : "EXPANDED.ALL", value })
     },
@@ -52,7 +75,7 @@ export function connect<T extends PropTypes>(service: TreeViewService, normalize
       send({ type: value ? "NODE.DESELECT" : "SELECTED.CLEAR", value })
     },
     select(value) {
-      send({ type: value ? "NODE.SELECT" : "SELECTED.ALL", value })
+      send({ type: value ? "NODE.SELECT" : "SELECTED.ALL", value, isTrusted: false })
     },
     getVisibleNodes() {
       return computed("visibleNodes").map(({ node }) => node)
@@ -388,5 +411,27 @@ export function connect<T extends PropTypes>(service: TreeViewService, normalize
         "data-depth": nodeState.depth,
       })
     },
+
+    getItemCheckboxProps(props) {
+      const nodeState = getNodeState(props)
+      return normalize.element({
+        ...parts.itemCheckbox.attrs,
+        tabIndex: -1,
+        role: "checkbox",
+        "data-state":
+          nodeState.checked === true ? "checked" : nodeState.checked === false ? "unchecked" : "indeterminate",
+        "aria-checked": nodeState.checked === true ? "true" : nodeState.checked === false ? "false" : "mixed",
+        onClick(event) {
+          if (event.defaultPrevented) return
+          if (nodeState.disabled) return
+
+          send({ type: "CHECKED.TOGGLE", value: nodeState.value, isBranch: nodeState.isBranch })
+          event.stopPropagation()
+
+          const node = event.currentTarget.closest("[role=treeitem]") as HTMLElement | null
+          node?.focus({ preventScroll: true })
+        },
+      })
+    },
   }
 }

packages/machines/tree-view/src/tree-view.connect.ts

@@ -1,10 +1,13 @@
-import type { TreeNode, TreeSkipFn } from "@zag-js/collection"
-import { createGuards, createMachine, type Params } from "@zag-js/core"
+import type { TreeNode } from "@zag-js/collection"
+import { createGuards, createMachine } from "@zag-js/core"
 import { getByTypeahead } from "@zag-js/dom-query"
-import { add, addOrRemove, diff, ensure, first, isArray, isEqual, last, remove, toArray, uniq } from "@zag-js/utils"
+import { addOrRemove, diff, first, isArray, isEqual, last, remove, toArray, uniq } from "@zag-js/utils"
 import { collection } from "./tree-view.collection"
 import * as dom from "./tree-view.dom"
 import type { TreeLoadingStatusMap, TreeViewSchema } from "./tree-view.types"
+import { toggleBranchChecked } from "./utils/checked-state"
+import { expandBranches } from "./utils/expand-branch"
+import { skipFn } from "./utils/visit-skip"
 
 const { and } = createGuards<TreeViewSchema>()
 
@@ -56,6 +59,14 @@ export const machine = createMachine<TreeViewSchema>({
       loadingStatus: bindable<TreeLoadingStatusMap>(() => ({
         defaultValue: {},
       })),
+      checkedValue: bindable(() => ({
+        defaultValue: prop("defaultCheckedValue") || [],
+        value: prop("checkedValue"),
+        isEqual,
+        onChange(value) {
+          prop("onCheckedChange")?.({ checkedValue: value })
+        },
+      })),
     }
   },
 
@@ -119,6 +130,15 @@ export const machine = createMachine<TreeViewSchema>({
     "NODE.DESELECT": {
       actions: ["deselectNode"],
     },
+    "CHECKED.TOGGLE": {
+      actions: ["toggleChecked"],
+    },
+    "CHECKED.SET": {
+      actions: ["setChecked"],
+    },
+    "CHECKED.CLEAR": {
+      actions: ["clearChecked"],
+    },
   },
 
   exit: ["clearPendingAborts"],
@@ -241,14 +261,12 @@ export const machine = createMachine<TreeViewSchema>({
       expandOnClick: ({ prop }) => !!prop("expandOnClick"),
     },
     actions: {
-      selectNode({ context, event, prop }) {
-        const value = toArray(event.id || event.value)
+      selectNode({ context, event }) {
+        const value = (event.id || event.value) as string | string[] | undefined
         context.set("selectedValue", (prev) => {
-          if (prop("selectionMode") === "single") {
-            return [last(value)].filter(Boolean)
-          } else {
-            return uniq([...prev, ...value])
-          }
+          if (value == null) return prev
+          if (!event.isTrusted && isArray(value)) return prev.concat(...value)
+          return [isArray(value) ? last(value) : value].filter(Boolean) as string[]
         })
       },
       deselectNode({ context, event }) {
@@ -480,114 +498,18 @@ export const machine = createMachine<TreeViewSchema>({
         aborts.forEach((abort) => abort.abort())
         aborts.clear()
       },
+      toggleChecked({ context, event, prop }) {
+        const collection = prop("collection")
+        context.set("checkedValue", (prev) =>
+          event.isBranch ? toggleBranchChecked(collection, event.value, prev) : addOrRemove(prev, event.value),
+        )
+      },
+      setChecked({ context, event }) {
+        context.set("checkedValue", event.value)
+      },
+      clearChecked({ context }) {
+        context.set("checkedValue", [])
+      },
     },
   },
 })
-
-function skipFn(params: Pick<Params<TreeViewSchema>, "prop" | "context">): TreeSkipFn<any> {
-  const { prop, context } = params
-  return function skip({ indexPath }) {
-    const paths = prop("collection").getValuePath(indexPath).slice(0, -1)
-    return paths.some((value) => !context.get("expandedValue").includes(value))
-  }
-}
-
-function partition<T>(array: T[], predicate: (value: T) => boolean) {
-  const pass: T[] = []
-  const fail: T[] = []
-  array.forEach((value) => {
-    if (predicate(value)) pass.push(value)
-    else fail.push(value)
-  })
-  return [pass, fail]
-}
-
-function expandBranches(params: Params<TreeViewSchema>, ids: string[]) {
-  const { context, prop, refs } = params
-
-  if (!prop("loadChildren")) {
-    context.set("expandedValue", (prev) => uniq(add(prev, ...ids)))
-    return
-  }
-
-  const [loadedValues, loadingValues] = partition(ids, (id) => context.get("loadingStatus")[id] === "loaded")
-
-  if (loadedValues.length > 0) {
-    context.set("expandedValue", (prev) => uniq(add(prev, ...loadedValues)))
-  }
-
-  if (loadingValues.length === 0) return
-
-  const collection = prop("collection")
-  const [nodeWithChildren, nodeWithoutChildren] = partition(loadingValues, (id) => {
-    const node = collection.findNode(id)
-    return collection.getNodeChildren(node).length > 0
-  })
-
-  // Check if node already has children (skip loading)
-  if (nodeWithChildren.length > 0) {
-    context.set("expandedValue", (prev) => uniq(add(prev, ...nodeWithChildren)))
-  }
-
-  if (nodeWithoutChildren.length === 0) return
-
-  context.set("loadingStatus", (prev) => ({
-    ...prev,
-    ...nodeWithoutChildren.reduce((acc, id) => ({ ...acc, [id]: "loading" }), {}),
-  }))
-
-  const nodesToLoad = nodeWithoutChildren.map((id) => {
-    const indexPath = collection.getIndexPath(id)!
-    const valuePath = collection.getValuePath(indexPath)!
-    const node = collection.findNode(id)!
-    return { id, indexPath, valuePath, node }
-  })
-
-  const pendingAborts = refs.get("pendingAborts")
-
-  // load children asynchronously
-  const loadChildren = prop("loadChildren")
-  ensure(loadChildren, () => "[zag-js/tree-view] `loadChildren` is required for async expansion")
-
-  const proms = nodesToLoad.map(({ id, indexPath, valuePath, node }) => {
-    const existingAbort = pendingAborts.get(id)
-    if (existingAbort) {
-      existingAbort.abort()
-      pendingAborts.delete(id)
-    }
-    const abortController = new AbortController()
-    pendingAborts.set(id, abortController)
-    return loadChildren({
-      valuePath,
-      indexPath,
-      node,
-      signal: abortController.signal,
-    })
-  })
-
-  // prefer `Promise.allSettled` over `Promise.all` to avoid early termination
-  Promise.allSettled(proms).then((results) => {
-    const loadedValues: string[] = []
-    const nextLoadingStatus = context.get("loadingStatus")
-    let collection = prop("collection")
-
-    results.forEach((result, index) => {
-      const { id, indexPath, node } = nodesToLoad[index]
-      if (result.status === "fulfilled") {
-        nextLoadingStatus[id] = "loaded"
-        loadedValues.push(id)
-        collection = collection.replace(indexPath, { ...node, children: result.value })
-      } else {
-        pendingAborts.delete(id)
-        Reflect.deleteProperty(nextLoadingStatus, id)
-      }
-    })
-
-    context.set("loadingStatus", nextLoadingStatus)
-
-    if (loadedValues.length) {
-      context.set("expandedValue", (prev) => uniq(add(prev, ...loadedValues)))
-      prop("onLoadChildrenComplete")?.({ collection })
-    }
-  })
-}

packages/machines/tree-view/src/tree-view.machine.ts

@@ -14,12 +14,16 @@ export const props = createProps<TreeViewProps>()([
   "onExpandedChange",
   "onFocusChange",
   "onSelectionChange",
+  "checkedValue",
   "selectedValue",
   "selectionMode",
   "typeahead",
   "defaultExpandedValue",
   "defaultSelectedValue",
+  "defaultCheckedValue",
+  "onCheckedChange",
   "onLoadChildrenComplete",
+  "onLoadChildrenError",
   "loadChildren",
 ])