feat(zui): add zui references for schemas unknown at build time (#330)

zui/.vscode/launch.json

@@ -0,0 +1,16 @@
+{
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Debug Test File",
+      "type": "node",
+      "request": "launch",
+      "autoAttachChildProcesses": true,
+      "skipFiles": ["<node_internals>/**", "**/node_modules/**"],
+      "program": "${workspaceRoot}/node_modules/vitest/vitest.mjs",
+      "args": ["run", "${relativeFile}"],
+      "smartStep": true,
+      "console": "integratedTerminal"
+    }
+  ]
+}

zui/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@bpinternal/zui",
-  "version": "0.8.6",
+  "version": "0.8.7",
   "description": "An extension of Zod for working nicely with UIs and JSON Schemas",
   "type": "module",
   "source": "./src/index.ts",

zui/src/deref.test.ts

@@ -0,0 +1,207 @@
+import { describe, test } from 'vitest'
+import { z } from './z/index'
+
+const foo = z.ref('foo')
+const bar = z.ref('bar')
+const baz = z.ref('baz')
+
+const deref = {
+  foo: z.string(),
+  bar: z.number(),
+  baz: z.boolean(),
+}
+
+const intersect = (...schemas: z.ZodTypeAny[]) => {
+  if (schemas.length === 0) {
+    throw new Error('Intersection expects at least one schema')
+  }
+
+  let current = schemas[0]!
+  for (let i = 1; i < schemas.length; i++) {
+    current = z.intersection(current, schemas[i]!)
+  }
+
+  return current
+}
+
+describe('dereference', () => {
+  test('array', () => {
+    const refSchema = z.array(bar)
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse([1, 2, 3])
+    expect(result.success).toBe(true)
+  })
+  test('discriminatedUnion', () => {
+    const refSchema = z.discriminatedUnion('type', [
+      z.object({ type: z.literal('foo'), foo: foo }),
+      z.object({ type: z.literal('bar'), bar: bar }),
+      z.object({ type: z.literal('baz'), baz: baz }),
+    ])
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse({ type: 'foo', foo: 'astring' })
+    expect(result.success).toBe(true)
+  })
+  test('function', () => {
+    const refSchema = z.function(z.tuple([foo, bar]), baz)
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse((_a: string, _b: number) => true)
+    expect(result.success).toBe(true)
+  })
+  test('intersection', () => {
+    const refSchema = intersect(z.object({ foo }), z.object({ bar }), z.object({ baz }))
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse({ foo: 'astring', bar: 1, baz: true })
+    expect(result.success).toBe(true)
+  })
+  test('map', () => {
+    const refSchema = z.map(foo, bar)
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse(new Map([['astring', 1]]))
+    expect(result.success).toBe(true)
+  })
+  test('nullable', () => {
+    const refSchema = z.nullable(foo)
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse(null)
+    expect(result.success).toBe(true)
+  })
+  test('object', () => {
+    const refSchema = z.object({
+      foo,
+      bar,
+      baz,
+    })
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse({ foo: 'astring', bar: 1, baz: true })
+    expect(result.success).toBe(true)
+  })
+  test('optional', () => {
+    const refSchema = z.optional(foo)
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse(undefined)
+    expect(result.success).toBe(true)
+  })
+  test('promise', () => {
+    const refSchema = z.promise(foo)
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse(Promise.resolve('astring'))
+    expect(result.success).toBe(true)
+  })
+  test('record', () => {
+    const refSchema = z.record(foo, bar)
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse({ foo: 1 })
+    expect(result.success).toBe(true)
+  })
+  test('set', () => {
+    const refSchema = z.set(foo)
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse(new Set(['astring']))
+    expect(result.success).toBe(true)
+  })
+  test('transformer', () => {
+    const refSchema = z.transformer(foo, {
+      type: 'transform',
+      transform: (data) => data,
+    })
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse('astring')
+    expect(result.success).toBe(true)
+  })
+  test('tuple', () => {
+    const refSchema = z.tuple([foo, bar])
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse(['astring', 1])
+    expect(result.success).toBe(true)
+  })
+  test('union', () => {
+    const refSchema = z.union([foo, bar, baz])
+    const derefSchema = refSchema.dereference(deref)
+    const result = derefSchema.safeParse('astring')
+    expect(result.success).toBe(true)
+  })
+})
+
+describe('getRegerences', () => {
+  test('array', () => {
+    const refSchema = z.array(bar)
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['bar'])
+  })
+  test('discriminatedUnion', () => {
+    const refSchema = z.discriminatedUnion('type', [
+      z.object({ type: z.literal('foo'), foo: foo }),
+      z.object({ type: z.literal('bar'), bar: bar }),
+      z.object({ type: z.literal('baz'), baz: baz }),
+    ])
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo', 'bar', 'baz'])
+  })
+  test('function', () => {
+    const refSchema = z.function(z.tuple([foo, bar]), baz)
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo', 'bar', 'baz'])
+  })
+  test('intersection', () => {
+    const refSchema = intersect(z.object({ foo }), z.object({ bar }), z.object({ baz }))
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo', 'bar', 'baz'])
+  })
+  test('map', () => {
+    const refSchema = z.map(foo, bar)
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo', 'bar'])
+  })
+  test('nullable', () => {
+    const refSchema = z.nullable(foo)
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo'])
+  })
+  test('object', () => {
+    const refSchema = z.object({
+      foo,
+      bar,
+      baz,
+    })
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo', 'bar', 'baz'])
+  })
+  test('optional', () => {
+    const refSchema = z.optional(foo)
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo'])
+  })
+  test('promise', () => {
+    const refSchema = z.promise(foo)
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo'])
+  })
+  test('record', () => {
+    const refSchema = z.record(foo, bar)
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo', 'bar'])
+  })
+  test('set', () => {
+    const refSchema = z.set(foo)
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo'])
+  })
+  test('transformer', () => {
+    const refSchema = z.transformer(foo, {
+      type: 'transform',
+      transform: (data) => data,
+    })
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo'])
+  })
+  test('tuple', () => {
+    const refSchema = z.tuple([foo, bar])
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo', 'bar'])
+  })
+  test('union', () => {
+    const refSchema = z.union([foo, bar, baz])
+    const refs = refSchema.getReferences()
+    expect(refs).toEqual(['foo', 'bar', 'baz'])
+  })
+})

zui/src/transforms/zui-to-json-schema/parseDef.ts

@@ -31,8 +31,8 @@ import { JsonSchema7UnknownType, parseUnknownDef } from './parsers/unknown'
 import { Refs, Seen } from './Refs'
 import { parseReadonlyDef } from './parsers/readonly'
 import { zuiKey } from '../../ui/constants'
+import { JsonSchema7RefType, parseRefDef } from './parsers/ref'
 
-type JsonSchema7RefType = { $ref: string }
 type JsonSchema7Meta = {
   default?: any
   description?: string
@@ -164,6 +164,8 @@ const selectParser = (def: any, typeName: ZodFirstPartyTypeKind, refs: Refs): Js
       return parseTupleDef(def, refs)
     case ZodFirstPartyTypeKind.ZodRecord:
       return parseRecordDef(def, refs)
+    case ZodFirstPartyTypeKind.ZodRef:
+      return parseRefDef(def)
     case ZodFirstPartyTypeKind.ZodLiteral:
       return parseLiteralDef(def, refs)
     case ZodFirstPartyTypeKind.ZodEnum:

zui/src/transforms/zui-to-json-schema/parsers/ref.ts

@@ -0,0 +1,11 @@
+import { ZodRefDef } from '../../../z/types/ref'
+
+export type JsonSchema7RefType = {
+  $ref: string
+}
+
+export function parseRefDef(def: ZodRefDef): JsonSchema7RefType {
+  return {
+    $ref: def.uri,
+  }
+}

zui/src/transforms/zui-to-json-schema/zui-extension.test.ts

@@ -381,4 +381,43 @@ describe('zuiToJsonSchema', () => {
       }
     `)
   })
+
+  test('generic is transformed to a ref', () => {
+    const T = z.ref('T')
+    const TJsonSchema = zuiToJsonSchema(T)
+    expect(TJsonSchema).toMatchInlineSnapshot(`
+      {
+        "$ref": "T",
+        "${zuiKey}": {},
+      }
+    `)
+
+    const schema = z.object({
+      description: z.string(),
+      data: T,
+    })
+
+    const jsonSchema = zuiToJsonSchema(schema)
+    expect(jsonSchema).toMatchInlineSnapshot(`
+      {
+        "additionalProperties": false,
+        "properties": {
+          "data": {
+            "$ref": "T",
+            "${zuiKey}": {},
+          },
+          "description": {
+            "type": "string",
+            "${zuiKey}": {},
+          },
+        },
+        "required": [
+          "description",
+          "data",
+        ],
+        "type": "object",
+        "${zuiKey}": {},
+      }
+    `)
+  })
 })

zui/src/transforms/zui-to-typescript-next/index.ts

@@ -328,6 +328,10 @@ ${escapeString((schema as z.ZodLiteral<any>).value)}`.trim()
     case z.ZodFirstPartyTypeKind.ZodReadonly:
       return `readonly ${sUnwrapZod(def.innerType, newConfig)}`
 
+    case z.ZodFirstPartyTypeKind.ZodRef:
+      // TODO: should be represented as a type argument <T>
+      throw new Error('ZodRef cannot be transformed to TypeScript yet')
+
     case z.ZodFirstPartyTypeKind.ZodTemplateLiteral:
       const inner = def.parts
         .map((p) => {

zui/src/z/types/array/index.ts

@@ -35,6 +35,17 @@ export class ZodArray<T extends ZodTypeAny, Cardinality extends ArrayCardinality
   ZodArrayDef<T>,
   Cardinality extends 'atleastone' ? [T['_input'], ...T['_input'][]] : T['_input'][]
 > {
+  dereference(defs: Record<string, ZodTypeAny>): ZodTypeAny {
+    return new ZodArray({
+      ...this._def,
+      type: this._def.type.dereference(defs),
+    })
+  }
+
+  getReferences(): string[] {
+    return this._def.type.getReferences()
+  }
+
   _parse(input: ParseInput): ParseReturnType<this['_output']> {
     const { ctx, status } = this._processInputParams(input)
 

zui/src/z/types/basetype/index.ts

@@ -156,6 +156,16 @@ export abstract class ZodType<Output = any, Def extends ZodTypeDef = ZodTypeDef,
 
   abstract _parse(input: ParseInput): ParseReturnType<Output>
 
+  /** deeply replace all references in the schema */
+  dereference(_defs: Record<string, ZodTypeAny>): ZodTypeAny {
+    return this
+  }
+
+  /** deeply scans the schema to check if it contains references */
+  getReferences(): string[] {
+    return []
+  }
+
   _getType(input: ParseInput): string {
     return getParsedType(input.data)
   }

zui/src/z/types/defs.test.ts

@@ -45,6 +45,8 @@ test('first party switch', () => {
       break
     case z.ZodFirstPartyTypeKind.ZodRecord:
       break
+    case z.ZodFirstPartyTypeKind.ZodRef:
+      break
     case z.ZodFirstPartyTypeKind.ZodMap:
       break
     case z.ZodFirstPartyTypeKind.ZodSet:

zui/src/z/types/defs.ts

@@ -92,6 +92,7 @@ export enum ZodFirstPartyTypeKind {
   ZodIntersection = 'ZodIntersection',
   ZodTuple = 'ZodTuple',
   ZodRecord = 'ZodRecord',
+  ZodRef = 'ZodRef',
   ZodMap = 'ZodMap',
   ZodSet = 'ZodSet',
   ZodFunction = 'ZodFunction',