feat(zui): Add discriminated union rendering (#270)

zui/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@bpinternal/zui",
-  "version": "0.6.4",
+  "version": "0.6.5",
   "description": "An extension of Zod for working nicely with UIs and JSON Schemas",
   "type": "module",
   "source": "./src/index.ts",
@@ -60,7 +60,7 @@
     "vitest": "1.5.0"
   },
   "dependencies": {
-    "@apidevtools/json-schema-ref-parser": "^11.5.4",
+    "@apidevtools/json-schema-ref-parser": "^11.5.5",
     "lodash": "^4.17.21",
     "react": "^18.2.0"
   },

zui/src/ui/constants.ts

@@ -1 +1,2 @@
 export const zuiKey = 'x-zui' as const
+export const ROOT = '$_ZROOT_$'

zui/src/ui/extensions.test.ts

@@ -44,6 +44,7 @@ const testExtensions = {
       params: z.object({ label: z.string().optional() }),
     },
   },
+  discriminatedUnion: {},
 } as const satisfies UIComponentDefinitions
 
 describe('ZUI UI Extensions', () => {

zui/src/ui/index.tsx

@@ -13,8 +13,8 @@ import {
   ZuiReactArrayChildProps,
   DefaultComponentDefinitions,
 } from './types'
-import { zuiKey } from './constants'
-import React, { type FC, useMemo } from 'react'
+import { ROOT, zuiKey } from './constants'
+import React, { type FC, useMemo, useEffect } from 'react'
 import { FormDataProvider, getDefaultItemData, useFormData } from './providers/FormDataProvider'
 import { getPathData } from './providers/FormDataProvider'
 import { formatTitle } from './titleutils'
@@ -26,11 +26,23 @@ type ComponentMeta<Type extends BaseType = BaseType> = {
   params: any
 }
 
+export const getSchemaType = (schema: JSONSchema): BaseType => {
+  if (schema.anyOf?.length) {
+    const discriminator = resolveDiscriminator(schema.anyOf)
+    return discriminator ? 'discriminatedUnion' : 'object'
+  }
+  if (schema.type === 'integer') {
+    return 'number'
+  }
+
+  return schema.type
+}
+
 const resolveComponent = <Type extends BaseType>(
   components: ZuiComponentMap<any> | undefined,
   fieldSchema: JSONSchema,
 ): ComponentMeta<Type> | null => {
-  const type = fieldSchema.type as BaseType
+  const type = getSchemaType(fieldSchema)
   const uiDefinition = fieldSchema[zuiKey]?.displayAs || null
 
   if (!uiDefinition || !Array.isArray(uiDefinition) || uiDefinition.length < 2) {
@@ -67,6 +79,68 @@ const resolveComponent = <Type extends BaseType>(
   }
 }
 
+export const resolveDiscriminator = (anyOf: ObjectSchema['anyOf']) => {
+  const output = anyOf
+    ?.map((schema) => {
+      if (schema.type !== 'object') {
+        return null
+      }
+      return Object.entries(schema.properties)
+        .map(([key, def]) => {
+          if (def.type === 'string' && def.const?.length) {
+            return { key, value: def.const }
+          }
+          return null
+        })
+        .filter((v): v is { key: string; value: string } => !!v)
+    })
+    .flat()
+    .reduce(
+      (acc, data) => {
+        if (!data) {
+          return acc
+        }
+        const { key, value } = data
+        if (acc.key === null) {
+          acc.key = key
+        }
+        if (acc.key === key) {
+          acc.values.push(value)
+        }
+
+        return acc
+      },
+      { key: null as string | null, values: [] as string[] },
+    )
+
+  if (output?.key === null || !output?.values.length) {
+    return null
+  }
+  return output
+}
+
+export const resolveDiscriminatedSchema = (key: string | null, value: string | null, anyOf: ObjectSchema['anyOf']) => {
+  if (!anyOf?.length || !key || !value) {
+    return null
+  }
+  for (const schema of anyOf) {
+    if (schema.type !== 'object') {
+      continue
+    }
+    const discriminator = schema.properties[key]
+    if (discriminator?.type === 'string' && discriminator.const === value) {
+      return {
+        ...schema,
+        properties: {
+          ...schema.properties,
+          [key]: { ...discriminator, [zuiKey]: { hidden: true } },
+        },
+      } as ObjectSchema
+    }
+  }
+  return null
+}
+
 export type ZuiFormProps<UI extends UIComponentDefinitions = DefaultComponentDefinitions> = {
   schema: JSONSchema
   components: ZuiComponentMap<UI>
@@ -122,7 +196,7 @@ const FormElementRenderer: FC<FormRendererProps> = ({ components, fieldSchema, p
 
   const { Component: _component, type } = componentMeta
 
-  const pathString = path.length > 0 ? path.join('.') : 'root'
+  const pathString = path.length > 0 ? path.join('.') : ROOT
 
   const baseProps: Omit<ZuiReactComponentBaseProps<BaseType, string, any>, 'data' | 'isArrayChild'> = {
     type,
@@ -208,6 +282,54 @@ const FormElementRenderer: FC<FormRendererProps> = ({ components, fieldSchema, p
       </Component>
     )
   }
+
+  if (type === 'discriminatedUnion') {
+    const Component = _component as any as ZuiReactComponent<'discriminatedUnion', string, any>
+
+    const discriminator = useMemo(() => resolveDiscriminator(fieldSchema.anyOf), [fieldSchema.anyOf])
+    const discriminatorValue = discriminator?.key ? data?.[discriminator.key] : null
+    useEffect(() => {
+      if (discriminator?.key && discriminator?.values.length) {
+        handlePropertyChange(pathString, { [discriminator.key]: discriminator.values[0] })
+      }
+    }, [])
+    const props: Omit<ZuiReactComponentProps<'discriminatedUnion', string, any>, 'children'> = {
+      ...baseProps,
+      type,
+      schema: baseProps.schema as any as ObjectSchema,
+      data: data || {},
+      discriminatorKey: discriminator?.key || null,
+      discriminatorOptions: discriminator?.values || null,
+      discriminatorValue,
+      setDiscriminator: (disc: string) => {
+        if (!discriminator?.key) {
+          console.warn('No discriminator key found, cannot set discriminator')
+          return
+        }
+        handlePropertyChange(pathString, { [discriminator.key]: disc })
+      },
+      ...childProps,
+    }
+    const discriminatedSchema = useMemo(
+      () => resolveDiscriminatedSchema(discriminator?.key || null, discriminatorValue, fieldSchema.anyOf),
+      [fieldSchema.anyOf, data, discriminator?.key],
+    )
+
+    return (
+      <Component key={baseProps.scope} {...props} isArrayChild={props.isArrayChild as any}>
+        {discriminatedSchema && (
+          <FormElementRenderer
+            components={components}
+            fieldSchema={discriminatedSchema}
+            path={path}
+            required={required}
+            key={path.join('.')}
+            isArrayChild={false}
+          />
+        )}
+      </Component>
+    )
+  }
   const Component = _component as any as ZuiReactComponent<any, any>
 
   const props: ZuiReactControlComponentProps<'boolean' | 'number' | 'string', string, any> = {

zui/src/ui/providers/FormDataProvider.tsx

@@ -2,6 +2,7 @@ import { PropsWithChildren, createContext, useContext, useMemo } from 'react'
 import React from 'react'
 import { JSONSchema } from '../types'
 import { jsonSchemaToZui } from '../../transforms/json-schema-to-zui'
+import { ROOT } from '../constants'
 
 export type FormFieldContextProps = {
   formData: any
@@ -67,6 +68,9 @@ export const useFormData = () => {
 }
 
 export function setObjectPath(obj: any, path: string, data: any): any {
+  if (path === ROOT) {
+    return data
+  }
   const pathArray = path.split('.')
   const pathArrayLength = pathArray.length
   pathArray.reduce((current: any, key: string, index: number) => {

zui/src/ui/stories/ZuiForm.stories.tsx

@@ -35,10 +35,12 @@ const exampleExtensions = {
       params: z.object({}),
     },
   },
+  discriminatedUnion: {},
 } as const satisfies UIComponentDefinitions
 
 const exampleSchema = z
   .object({
+    root: z.string().title('Root').placeholder('Root').default('root'),
     field: z.enum(['yes', 'no']).displayAs<typeof exampleExtensions>({ id: 'debug', params: {} }),
     firstName: z.string().title('first name').disabled().placeholder('Enter your name').nullable(),
     lastName: z.string().min(3).title('Last Name <3').optional().nullable(),
@@ -54,7 +56,21 @@ const exampleSchema = z
       .nonempty(),
     // tests the hidden function
     aRandomField: z.string().optional().hidden(),
-
+    aDiscriminatedUnion: z.discriminatedUnion('type', [
+      z.object({ type: z.literal('text'), text: z.string().placeholder('Some text') }),
+      z.object({ type: z.literal('number'), b: z.number().placeholder('42').default(5) }),
+      z.object({
+        type: z.literal('complex'),
+        address: z.object({
+          street: z.string().placeholder('1234 Main St'),
+          city: z.string().placeholder('San Francisco'),
+          state: z.string().placeholder('CA'),
+          zip: z.string().placeholder('94111'),
+        }),
+        root: z.string().placeholder('root'),
+        babies: z.array(z.object({ name: z.string(), age: z.number() })),
+      }),
+    ]),
     stuff: z.object({
       birthday: z.string(),
       plan: z.enum(['basic', 'premium']),
@@ -155,6 +171,7 @@ const componentMap: ZuiComponentMap<typeof exampleExtensions> = {
             placeholder={zuiProps?.placeholder}
             onChange={(e) => onChange(parseFloat(e.target.value))}
             value={data || 0}
+            defaultValue={schema.default || 0}
           />
           {required && <span>*</span>}
           <ErrorBox errors={errors} data={data} />
@@ -204,6 +221,23 @@ const componentMap: ZuiComponentMap<typeof exampleExtensions> = {
       )
     },
   },
+  discriminatedUnion: {
+    default: ({ children, discriminatorKey, discriminatorOptions, discriminatorValue, setDiscriminator }) => {
+      return (
+        <div>
+          <span>{discriminatorKey}</span>
+          <select value={discriminatorValue || undefined} onChange={(e) => setDiscriminator(e.target.value)}>
+            {discriminatorOptions?.map((option) => (
+              <option key={option} value={option}>
+                {option}
+              </option>
+            ))}
+          </select>
+          {children}
+        </div>
+      )
+    },
+  },
 }
 
 const ZuiFormExample = () => {
@@ -214,9 +248,7 @@ const ZuiFormExample = () => {
   return (
     <>
       <ZuiForm<typeof exampleExtensions>
-        schema={exampleSchema.toJsonSchema({
-          target: 'jsonSchema7',
-        })}
+        schema={exampleSchema.toJsonSchema()}
         value={formData}
         onChange={setFormData}
         components={componentMap}

zui/src/ui/types.ts

@@ -4,6 +4,9 @@ import { zuiKey } from './constants'
 
 export type BaseSchema = {
   description?: string
+  anyOf?: JSONSchema[]
+  oneOf?: JSONSchema[]
+  allOf?: JSONSchema[]
   [zuiKey]?: {
     tooltip?: boolean
     disabled?: boolean
@@ -84,11 +87,13 @@ export type JSONSchemaOfType<T extends BaseType> = T extends 'string'
         ? ObjectSchema
         : T extends 'array'
           ? ArraySchema
-          : never
+          : T extends 'discriminatedUnion'
+            ? ObjectSchema
+            : never
 
-export type BaseType = 'number' | 'string' | 'boolean' | 'object' | 'array'
+export type BaseType = 'number' | 'string' | 'boolean' | 'object' | 'array' | 'discriminatedUnion'
 
-export const containerTypes = ['object', 'array'] as const
+export const containerTypes = ['object', 'array', 'discriminatedUnion'] as const
 export type ContainerType = (typeof containerTypes)[number]
 
 export type DefaultComponentDefinitions = {
@@ -97,6 +102,7 @@ export type DefaultComponentDefinitions = {
   boolean: {}
   object: {}
   array: {}
+  discriminatedUnion: {}
 }
 
 export type UIComponentDefinitions = {
@@ -125,7 +131,12 @@ export type ZodKindToBaseType<T extends z.ZodTypeDef> = T extends infer U
                 ? ZodKindToBaseType<U['innerType']['_def']>
                 : U extends { typeName: z.ZodFirstPartyTypeKind.ZodNullable; innerType: z.ZodTypeAny }
                   ? ZodKindToBaseType<U['innerType']['_def']>
-                  : never
+                  : U extends {
+                        typeName: z.ZodFirstPartyTypeKind.ZodDiscriminatedUnion
+                        options: z.ZodDiscriminatedUnionOption<any>[]
+                      }
+                    ? 'discriminatedUnion'
+                    : never
   : never
 
 export type BaseTypeToType<T extends BaseType> = T extends 'string'
@@ -195,6 +206,18 @@ export type ZuiReactArrayChildProps =
       isArrayChild: false
     }
 
+export type ZuiReactDiscriminatedUnionComponentProps<
+  Type extends ContainerType,
+  ID extends keyof UI[Type],
+  UI extends UIComponentDefinitions = DefaultComponentDefinitions,
+> = ZuiReactComponentBaseProps<Type, ID, UI> & {
+  discriminatorKey: string | null
+  discriminatorOptions: string[] | null
+  discriminatorValue: string | null
+  setDiscriminator: (discriminator: string) => void
+  children: JSX.Element | JSX.Element[] | null
+}
+
 export type ZuiReactObjectComponentProps<
   Type extends ContainerType,
   ID extends keyof UI[Type],
@@ -231,7 +254,9 @@ export type ZuiReactComponentProps<
   ? ZuiReactObjectComponentProps<Type, ID, UI>
   : Type extends 'array'
     ? ZuiReactArrayComponentProps<Type, ID, UI>
-    : ZuiReactControlComponentProps<Type, ID, UI>
+    : Type extends 'discriminatedUnion'
+      ? ZuiReactDiscriminatedUnionComponentProps<Type, ID, UI>
+      : ZuiReactControlComponentProps<Type, ID, UI>
 
 export type ZuiReactComponent<
   Type extends BaseType,