feat(context): improve createContext type definitions

packages/react/src/components/checkbox/use-checkbox-group-context.tsx

@@ -3,9 +3,7 @@ import type { UseCheckboxGroupReturn } from './use-checkbox-group'
 
 export interface UseCheckboxGroupContext extends UseCheckboxGroupReturn {}
 
-export const [CheckboxGroupContextProvider, useCheckboxGroupContext] = createContext<
-  UseCheckboxGroupContext | undefined
->({
+export const [CheckboxGroupContextProvider, useCheckboxGroupContext] = createContext<UseCheckboxGroupContext>({
   name: 'CheckboxGroupContext',
   hookName: 'useCheckboxGroupContext',
   providerName: '<CheckboxGroupProvider />',

packages/react/src/components/field/examples/custom-control.tsx

@@ -4,7 +4,7 @@ import styles from 'styles/field.module.css'
 export const CustomControl = () => (
   <Field.Root className={styles.Root} invalid>
     <Field.Label className={styles.Label}>Any Control</Field.Label>
-    <Field.Context>{(context) => <input {...context.getInputProps()} />}</Field.Context>
+    <Field.Context>{(context) => <input {...context?.getInputProps()} />}</Field.Context>
     <Field.HelperText className={styles.HelperText}>Uses getInputProps() for maximum flexibility</Field.HelperText>
     <Field.ErrorText className={styles.ErrorText}>This field has an error</Field.ErrorText>
   </Field.Root>

packages/react/src/components/field/field-context.tsx

@@ -1,8 +1,8 @@
 import type { ReactNode } from 'react'
-import { type UseFieldContext, useFieldContext } from './use-field-context'
+import { useFieldContext } from './use-field-context'
 
 export interface FieldContextProps {
-  children: (context: UseFieldContext) => ReactNode
+  children: (context: ReturnType<typeof useFieldContext>) => ReactNode
 }
 
 export const FieldContext = (props: FieldContextProps) => props.children(useFieldContext())

packages/react/src/components/field/field-error-text.tsx

@@ -8,7 +8,7 @@ export interface FieldErrorTextProps extends HTMLProps<'span'>, FieldErrorTextBa
 
 export const FieldErrorText = forwardRef<HTMLSpanElement, FieldErrorTextProps>((props, ref) => {
   const field = useFieldContext()
-  const mergedProps = mergeProps(field.getErrorTextProps(), props)
+  const mergedProps = mergeProps(field?.getErrorTextProps(), props)
 
   if (field?.invalid) {
     return <ark.span {...mergedProps} ref={ref} />

packages/react/src/components/field/field-required-indicator.tsx

@@ -12,11 +12,11 @@ export const FieldRequiredIndicator = forwardRef<HTMLSpanElement, FieldRequiredI
   ({ fallback, ...props }, ref) => {
     const field = useFieldContext()
 
-    if (!field.required) {
+    if (!field?.required) {
       return fallback
     }
 
-    const mergedProps = mergeProps(field.getRequiredIndicatorProps(), props)
+    const mergedProps = mergeProps(field?.getRequiredIndicatorProps(), props)
     return (
       <ark.span {...mergedProps} ref={ref}>
         {props.children ?? '*'}

packages/react/src/components/fieldset/fieldset-context.tsx

@@ -1,8 +1,8 @@
 import type { ReactNode } from 'react'
-import { type UseFieldsetContext, useFieldsetContext } from './use-fieldset-context'
+import { useFieldsetContext } from './use-fieldset-context'
 
 export interface FieldsetContextProps {
-  children: (context: UseFieldsetContext) => ReactNode
+  children: (context: ReturnType<typeof useFieldsetContext>) => ReactNode
 }
 
 export const FieldsetContext = (props: FieldsetContextProps) => props.children(useFieldsetContext())

packages/react/src/components/fieldset/fieldset-error-text.tsx

@@ -8,9 +8,9 @@ export interface FieldsetErrorTextProps extends HTMLProps<'span'>, FieldsetError
 
 export const FieldsetErrorText = forwardRef<HTMLSpanElement, FieldsetErrorTextProps>((props, ref) => {
   const fieldset = useFieldsetContext()
-  const mergedProps = mergeProps(fieldset.getErrorTextProps(), props)
+  const mergedProps = mergeProps(fieldset?.getErrorTextProps(), props)
 
-  return fieldset.invalid ? <ark.span {...mergedProps} ref={ref} /> : null
+  return fieldset?.invalid ? <ark.span {...mergedProps} ref={ref} /> : null
 })
 
 FieldsetErrorText.displayName = 'FieldsetErrorText'

packages/react/src/components/fieldset/fieldset-helper-text.tsx

@@ -8,7 +8,7 @@ export interface FieldsetHelperTextProps extends HTMLProps<'span'>, FieldsetHelp
 
 export const FieldsetHelperText = forwardRef<HTMLSpanElement, FieldsetHelperTextProps>((props, ref) => {
   const fieldset = useFieldsetContext()
-  const mergedProps = mergeProps(fieldset.getHelperTextProps(), props)
+  const mergedProps = mergeProps(fieldset?.getHelperTextProps(), props)
 
   return <ark.span {...mergedProps} ref={ref} />
 })

packages/react/src/components/fieldset/fieldset-legend.tsx

@@ -8,7 +8,7 @@ export interface FieldsetLegendProps extends HTMLProps<'legend'>, FieldsetLegend
 
 export const FieldsetLegend = forwardRef<HTMLLegendElement, FieldsetLegendProps>((props, ref) => {
   const fieldset = useFieldsetContext()
-  const mergedProps = mergeProps(fieldset.getLegendProps(), props)
+  const mergedProps = mergeProps(fieldset?.getLegendProps(), props)
 
   return <ark.legend {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-arrow-tip.tsx

@@ -8,7 +8,7 @@ export interface MenuArrowTipProps extends HTMLProps<'div'>, MenuArrowTipBasePro
 
 export const MenuArrowTip = forwardRef<HTMLDivElement, MenuArrowTipProps>((props, ref) => {
   const menu = useMenuContext()
-  const mergedProps = mergeProps(menu.getArrowTipProps(), props)
+  const mergedProps = mergeProps(menu?.getArrowTipProps(), props)
 
   return <ark.div {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-arrow.tsx

@@ -8,7 +8,7 @@ export interface MenuArrowProps extends HTMLProps<'div'>, MenuArrowBaseProps {}
 
 export const MenuArrow = forwardRef<HTMLDivElement, MenuArrowProps>((props, ref) => {
   const menu = useMenuContext()
-  const mergedProps = mergeProps(menu.getArrowProps(), props)
+  const mergedProps = mergeProps(menu?.getArrowProps(), props)
 
   return <ark.div {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-checkbox-item.tsx

@@ -28,12 +28,12 @@ export const MenuCheckboxItem = forwardRef<HTMLDivElement, MenuCheckboxItemProps
     type: 'checkbox',
   }
   const menu = useMenuContext()
-  const mergedProps = mergeProps(menu.getOptionItemProps(optionItemProps), localProps)
-  const optionItemState = menu.getOptionItemState(optionItemProps)
+  const mergedProps = mergeProps(menu?.getOptionItemProps(optionItemProps), localProps)
+  const optionItemState = menu?.getOptionItemState(optionItemProps)
 
   return (
     <MenuItemPropsProvider value={optionItemProps}>
-      <MenuItemProvider value={optionItemState}>
+      <MenuItemProvider value={optionItemState!}>
         <ark.div {...mergedProps} ref={ref} />
       </MenuItemProvider>
     </MenuItemPropsProvider>

packages/react/src/components/menu/menu-content.tsx

@@ -11,7 +11,7 @@ export interface MenuContentProps extends HTMLProps<'div'>, MenuContentBaseProps
 export const MenuContent = forwardRef<HTMLDivElement, MenuContentProps>((props, ref) => {
   const menu = useMenuContext()
   const presence = usePresenceContext()
-  const mergedProps = mergeProps(menu.getContentProps(), presence.getPresenceProps(), props)
+  const mergedProps = mergeProps(menu?.getContentProps(), presence.getPresenceProps(), props)
 
   if (presence.unmounted) {
     return null

packages/react/src/components/menu/menu-context-trigger.tsx

@@ -8,7 +8,7 @@ export interface MenuContextTriggerProps extends HTMLProps<'button'>, MenuContex
 
 export const MenuContextTrigger = forwardRef<HTMLButtonElement, MenuContextTriggerProps>((props, ref) => {
   const menu = useMenuContext()
-  const mergedProps = mergeProps(menu.getContextTriggerProps(), props)
+  const mergedProps = mergeProps(menu?.getContextTriggerProps(), props)
 
   return <ark.button {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-context.tsx

@@ -1,8 +1,8 @@
 import type { ReactNode } from 'react'
-import { type UseMenuContext, useMenuContext } from './use-menu-context'
+import { useMenuContext } from './use-menu-context'
 
 export interface MenuContextProps {
-  children: (context: UseMenuContext) => ReactNode
+  children: (context: ReturnType<typeof useMenuContext>) => ReactNode
 }
 
 export const MenuContext = (props: MenuContextProps) => props.children(useMenuContext())

packages/react/src/components/menu/menu-indicator.tsx

@@ -8,7 +8,7 @@ export interface MenuIndicatorProps extends HTMLProps<'div'>, MenuIndicatorBaseP
 
 export const MenuIndicator = forwardRef<HTMLDivElement, MenuIndicatorProps>((props, ref) => {
   const menu = useMenuContext()
-  const mergedProps = mergeProps(menu.getIndicatorProps(), props)
+  const mergedProps = mergeProps(menu?.getIndicatorProps(), props)
 
   return <ark.div {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-item-group-label.tsx

@@ -10,7 +10,7 @@ export interface MenuItemGroupLabelProps extends HTMLProps<'div'>, MenuItemGroup
 export const MenuItemGroupLabel = forwardRef<HTMLDivElement, MenuItemGroupLabelProps>((props, ref) => {
   const menu = useMenuContext()
   const itemGroup = useMenuItemGroupContext()
-  const mergedProps = mergeProps(menu.getItemGroupLabelProps({ htmlFor: itemGroup.id }), props)
+  const mergedProps = mergeProps(menu?.getItemGroupLabelProps({ htmlFor: itemGroup.id }), props)
 
   return <ark.div {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-item-group.tsx

@@ -19,7 +19,7 @@ export const MenuItemGroup = forwardRef<HTMLDivElement, MenuItemGroupProps>((pro
   const menu = useMenuContext()
   const id = useId()
   const itemGroupProps = { id, ...optionalItemGroupProps }
-  const mergedProps = mergeProps(menu.getItemGroupProps(itemGroupProps), localProps)
+  const mergedProps = mergeProps(menu?.getItemGroupProps(itemGroupProps), localProps)
 
   return (
     <MenuItemGroupProvider value={itemGroupProps}>

packages/react/src/components/menu/menu-item-indicator.tsx

@@ -10,7 +10,7 @@ export interface MenuItemIndicatorProps extends HTMLProps<'div'>, MenuItemIndica
 export const MenuItemIndicator = forwardRef<HTMLDivElement, MenuItemIndicatorProps>((props, ref) => {
   const menu = useMenuContext()
   const itemProps = useMenuItemPropsContext()
-  const mergedProps = mergeProps(menu.getItemIndicatorProps(itemProps), props)
+  const mergedProps = mergeProps(menu?.getItemIndicatorProps(itemProps), props)
 
   return <ark.div {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-item-text.tsx

@@ -10,7 +10,7 @@ export interface MenuItemTextProps extends HTMLProps<'div'>, MenuItemTextBasePro
 export const MenuItemText = forwardRef<HTMLDivElement, MenuItemTextProps>((props, ref) => {
   const menu = useMenuContext()
   const itemProps = useMenuItemPropsContext()
-  const mergedProps = mergeProps(menu.getItemTextProps(itemProps), props)
+  const mergedProps = mergeProps(menu?.getItemTextProps(itemProps), props)
 
   return <ark.div {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-item.tsx

@@ -31,17 +31,18 @@ export const MenuItem = forwardRef<HTMLDivElement, MenuItemProps>((props, ref) =
   ])
 
   const menu = useMenuContext()
-  const mergedProps = mergeProps(menu.getItemProps(itemProps), localProps)
-  const itemState = menu.getItemState(itemProps)
+  const mergedProps = mergeProps(menu?.getItemProps(itemProps), localProps)
+  const itemState = menu?.getItemState(itemProps)
 
   // biome-ignore lint/correctness/useExhaustiveDependencies: intentional
   useEffect(() => {
-    return menu.addItemListener({ id: itemState.id, onSelect: itemProps.onSelect })
-  }, [itemState.id, itemProps.onSelect])
+    if (!itemState?.id) return
+    return menu?.addItemListener({ id: itemState.id, onSelect: itemProps.onSelect })
+  }, [itemState?.id, itemProps.onSelect])
 
   return (
     <MenuItemPropsProvider value={itemProps}>
-      <MenuItemProvider value={itemState}>
+      <MenuItemProvider value={itemState!}>
         <ark.div {...mergedProps} ref={ref} />
       </MenuItemProvider>
     </MenuItemPropsProvider>

packages/react/src/components/menu/menu-positioner.tsx

@@ -9,7 +9,7 @@ export interface MenuPositionerProps extends HTMLProps<'div'>, MenuPositionerBas
 
 export const MenuPositioner = forwardRef<HTMLDivElement, MenuPositionerProps>((props, ref) => {
   const menu = useMenuContext()
-  const mergedProps = mergeProps(menu.getPositionerProps(), props)
+  const mergedProps = mergeProps(menu?.getPositionerProps(), props)
   const presence = usePresenceContext()
 
   if (presence.unmounted) {

packages/react/src/components/menu/menu-radio-item-group.tsx

@@ -18,7 +18,7 @@ export const MenuRadioItemGroup = forwardRef<HTMLDivElement, MenuRadioItemGroupP
   const menu = useMenuContext()
   const id = useId()
   const itemGroupProps = { id, ...optionalItemGroupProps }
-  const mergedProps = mergeProps(menu.getItemGroupProps({ id: itemGroupProps.id }), localProps)
+  const mergedProps = mergeProps(menu?.getItemGroupProps({ id: itemGroupProps.id }), localProps)
 
   return (
     <MenuItemGroupProvider value={itemGroupProps}>

packages/react/src/components/menu/menu-radio-item.tsx

@@ -30,12 +30,12 @@ export const MenuRadioItem = forwardRef<HTMLDivElement, MenuRadioItemProps>((pro
     type: 'radio',
     onCheckedChange: () => itemGroup.onValueChange?.({ value: partialItemProps.value }),
   }
-  const mergedProps = mergeProps(menu.getOptionItemProps(optionItemProps), localProps)
-  const optionItemState = menu.getOptionItemState(optionItemProps)
+  const mergedProps = mergeProps(menu?.getOptionItemProps(optionItemProps), localProps)
+  const optionItemState = menu?.getOptionItemState(optionItemProps)
 
   return (
     <MenuItemPropsProvider value={optionItemProps}>
-      <MenuItemProvider value={optionItemState}>
+      <MenuItemProvider value={optionItemState!}>
         <ark.div {...mergedProps} ref={ref} />
       </MenuItemProvider>
     </MenuItemPropsProvider>

packages/react/src/components/menu/menu-separator.tsx

@@ -8,7 +8,7 @@ export interface MenuSeparatorProps extends HTMLProps<'hr'>, MenuSeparatorBasePr
 
 export const MenuSeparator = forwardRef<HTMLHRElement, MenuSeparatorProps>((props, ref) => {
   const menu = useMenuContext()
-  const mergedProps = mergeProps(menu.getSeparatorProps(), props)
+  const mergedProps = mergeProps(menu?.getSeparatorProps(), props)
 
   return <ark.hr {...mergedProps} ref={ref} />
 })

packages/react/src/components/menu/menu-trigger.tsx

@@ -12,8 +12,8 @@ export const MenuTrigger = forwardRef<HTMLButtonElement, MenuTriggerProps>((prop
   const presence = usePresenceContext()
   const mergedProps = mergeProps(
     {
-      ...menu.getTriggerProps(),
-      'aria-controls': presence.unmounted ? undefined : menu.getTriggerProps()['aria-controls'],
+      ...menu?.getTriggerProps(),
+      'aria-controls': presence.unmounted ? undefined : menu?.getTriggerProps()['aria-controls'],
     },
     props,
   )

packages/react/src/components/navigation-menu/navigation-menu-content.tsx

@@ -21,7 +21,7 @@ export const NavigationMenuContent = forwardRef<HTMLDivElement, NavigationMenuCo
   const itemContext = useNavigationMenuItemPropsContext()
 
   const value = props.value ?? itemContext?.value
-  const [contentProps, localProps] = splitContentProps({ ...props, value }, ['value'])
+  const [contentProps, localProps] = splitContentProps({ ...props, value: value! }, ['value'])
 
   const renderStrategyProps = useRenderStrategyPropsContext()
   const presence = usePresence({ ...renderStrategyProps, present: api.value === value })

packages/react/src/components/navigation-menu/navigation-menu-item-indicator.tsx

@@ -11,7 +11,7 @@ export const NavigationMenuItemIndicator = forwardRef<HTMLDivElement, Navigation
   (props, ref) => {
     const navigationMenu = useNavigationMenuContext()
     const itemProps = useNavigationMenuItemPropsContext()
-    const mergedProps = mergeProps(navigationMenu.getItemIndicatorProps(itemProps), props)
+    const mergedProps = mergeProps(navigationMenu.getItemIndicatorProps(itemProps!), props)
 
     return <ark.div {...mergedProps} ref={ref} />
   },

packages/react/src/components/navigation-menu/navigation-menu-link.tsx

@@ -16,7 +16,12 @@ export const NavigationMenuLink = forwardRef<HTMLAnchorElement, NavigationMenuLi
   const itemContext = useNavigationMenuItemPropsContext()
   const value = props.value ?? itemContext?.value
 
-  const [linkProps, localProps] = splitLinkProps({ ...props, value }, ['current', 'onSelect', 'value', 'closeOnClick'])
+  const [linkProps, localProps] = splitLinkProps({ ...props, value: value! }, [
+    'current',
+    'onSelect',
+    'value',
+    'closeOnClick',
+  ])
   const navigationMenu = useNavigationMenuContext()
   const mergedProps = mergeProps(navigationMenu.getLinkProps(linkProps), localProps)
 

packages/react/src/utils/create-context.ts

@@ -1,6 +1,10 @@
 import { hasProp, isFunction } from '@zag-js/utils'
 import { createContext as createReactContext, useContext as useReactContext } from 'react'
 
+function getErrorMessage(hook: string, provider: string) {
+  return `${hook} returned \`undefined\`. Seems you forgot to wrap component within ${provider}`
+}
+
 interface CreateContextOptions<T> {
   strict?: boolean | undefined
   hookName?: string | undefined
@@ -10,11 +14,19 @@ interface CreateContextOptions<T> {
   defaultValue?: T | undefined
 }
 
-type CreateContextReturn<T> = [React.Provider<T>, () => T, React.Context<T>]
+type CreateContextReturn<T, HasValue extends boolean> = [
+  React.Provider<T>,
+  HasValue extends true ? () => T : () => T | undefined,
+  HasValue extends true ? React.Context<T> : React.Context<T | undefined>,
+]
 
-function getErrorMessage(hook: string, provider: string) {
-  return `${hook} returned \`undefined\`. Seems you forgot to wrap component within ${provider}`
-}
+export function createContext<T>(options?: CreateContextOptions<T> & { strict?: true }): CreateContextReturn<T, true>
+
+export function createContext<T>(
+  options: CreateContextOptions<T> & { strict: false; defaultValue: T },
+): CreateContextReturn<T, true>
+
+export function createContext<T>(options: CreateContextOptions<T> & { strict: false }): CreateContextReturn<T, false>
 
 export function createContext<T>(options: CreateContextOptions<T> = {}) {
   const {
@@ -45,5 +57,5 @@ export function createContext<T>(options: CreateContextOptions<T> = {}) {
     return context
   }
 
-  return [Context.Provider, useContext, Context] as CreateContextReturn<T>
+  return [Context.Provider, useContext, Context] as any
 }