@kubb/plugin-zod 5.0.0-beta.4 → 5.0.0-beta.56

package/src/printers/printerZod.ts

@@ -1,8 +1,8 @@
-import { stringify } from '@internals/utils'
+import { buildList, buildObject, extractRefName, objectKey, stringify } from '@kubb/ast/utils'
 
 import { ast } from '@kubb/core'
 import type { PluginZod, ResolverZod } from '../types.ts'
-import { applyModifiers, formatLiteral, lengthConstraints, numberConstraints, shouldCoerce } from '../utils.ts'
+import { applyModifiers, containsCodec, formatLiteral, getCodec, lengthConstraints, numberConstraints, shouldCoerce } from '../utils.ts'
 import type { AdapterOas } from '@kubb/adapter-oas'
 
 /**
@@ -53,12 +53,21 @@ export type PrinterZodOptions = {
   /**
    * Properties to exclude using `.omit({ key: true })`.
    */
-  keysToOmit?: Array<string>
+  keysToOmit?: Array<string> | null
   /**
    * Schema names that form circular dependency chains.
    * Properties referencing these emit lazy getters wrapping refs in `z.lazy(() => …)`.
    */
   cyclicSchemas?: ReadonlySet<string>
+  /**
+   * Print direction for `dateType: 'date'` fields (`Date` in TypeScript):
+   * - `'output'` (default) — decode the wire `string` into a `Date` (response bodies).
+   * - `'input'` — encode a `Date` back into the wire `string` (request bodies/params).
+   *
+   * Diverging the directions requires the generator to emit an `${name}InputSchema`
+   * variant for each date-bearing component.
+   */
+  direction?: 'input' | 'output'
   /**
    * Custom handler map for node type overrides.
    */
@@ -70,6 +79,33 @@ export type PrinterZodOptions = {
  */
 export type PrinterZodFactory = ast.PrinterFactoryOptions<'zod', PrinterZodOptions, string, string>
 
+function strictOneOfMember(member: string, node: ast.SchemaNode): string {
+  if (node.type === 'object' && node.additionalProperties === undefined) {
+    return `${member}.strict()`
+  }
+
+  if (node.type === 'ref') {
+    if (member.startsWith('z.lazy(')) {
+      return member
+    }
+
+    const schema = ast.syncSchemaRef(node)
+
+    if (schema.type === 'object' && (schema.additionalProperties === undefined || schema.additionalProperties === false)) {
+      return `${member}.strict()`
+    }
+  }
+
+  return member
+}
+
+function getMemberConstraint(member: ast.SchemaNode): string | undefined {
+  if (member.primitive === 'string') return lengthConstraints(ast.narrowSchema(member, 'string') ?? {}) || undefined
+  if (member.primitive === 'number' || member.primitive === 'integer')
+    return numberConstraints(ast.narrowSchema(member, 'number') ?? ast.narrowSchema(member, 'integer') ?? {}) || undefined
+  if (member.primitive === 'array') return lengthConstraints(ast.narrowSchema(member, 'array') ?? {}) || undefined
+}
+
 /**
  * Zod v4 printer built with `definePrinter`.
  *
@@ -112,11 +148,13 @@ export const printerZod = ast.definePrinter<PrinterZodFactory>((options) => {
         return shouldCoerce(this.options.coercion, 'numbers') ? 'z.coerce.bigint()' : 'z.bigint()'
       },
       date(node) {
-        if (node.representation === 'string') {
-          return 'z.iso.date()'
+        // representation: 'date' → typed as `Date`; decode/encode at the boundary.
+        const codec = getCodec(node)
+        if (codec) {
+          return this.options.direction === 'input' ? codec.encode(node) : codec.decode(node)
         }
 
-        return shouldCoerce(this.options.coercion, 'dates') ? 'z.coerce.date()' : 'z.date()'
+        return 'z.iso.date()'
       },
       datetime(node) {
         const offset = node.offset || this.options.dateType === 'stringOffset'
@@ -164,9 +202,17 @@ export const printerZod = ast.definePrinter<PrinterZodFactory>((options) => {
         return `z.enum([${nonNullValues.map(formatLiteral).join(', ')}])`
       },
       ref(node) {
-        if (!node.name) return undefined
-        const refName = node.ref ? (ast.extractRefName(node.ref) ?? node.name) : node.name
-        const resolvedName = node.ref ? (this.options.resolver?.default(refName, 'function') ?? refName) : node.name
+        if (!node.name) return null
+        const refName = node.ref ? (extractRefName(node.ref) ?? node.name) : node.name
+
+        // In the input direction, a date-bearing component resolves to its `${name}InputSchema`
+        // variant so request bodies encode `Date → string` instead of decoding.
+        const useInputVariant = node.ref != null && this.options.direction === 'input' && containsCodec(node)
+        const resolvedName = node.ref
+          ? useInputVariant
+            ? (this.options.resolver?.resolveInputSchemaName(refName) ?? refName)
+            : (this.options.resolver?.default(refName, 'function') ?? refName)
+          : node.name
 
         if (node.ref && this.options.cyclicSchemas?.has(refName)) {
           return `z.lazy(() => ${resolvedName})`
@@ -175,93 +221,92 @@ export const printerZod = ast.definePrinter<PrinterZodFactory>((options) => {
         return resolvedName
       },
       object(node) {
-        const properties = node.properties
-          .map((prop) => {
-            const { name: propName, schema } = prop
-
-            const meta = ast.syncSchemaRef(schema)
-
-            const isNullable = meta.nullable
-            const isOptional = schema.optional
-            const isNullish = schema.nullish
-
-            const hasSelfRef = this.options.cyclicSchemas != null && ast.containsCircularRef(schema, { circularSchemas: this.options.cyclicSchemas })
-            // Inside a getter the getter itself defers evaluation, so suppress
-            // z.lazy() wrapping on nested refs by temporarily clearing cyclicSchemas.
-            if (hasSelfRef) this.options.cyclicSchemas = undefined
-            const baseOutput = this.transform(schema) ?? this.transform(ast.createSchema({ type: 'unknown' }))!
-            if (hasSelfRef) this.options.cyclicSchemas = options.cyclicSchemas
-
-            const wrappedOutput = this.options.wrapOutput ? this.options.wrapOutput({ output: baseOutput, schema }) || baseOutput : baseOutput
-
-            // When a property schema is not a ref but the metadata is from a ref (e.g., discriminator
-            // property override), skip applying the description from the ref target to avoid applying
-            // metadata from a replaced schema.
-            let descriptionToApply = meta.description
-            if (schema.type !== 'ref' && meta.type === 'ref') {
-              descriptionToApply = undefined
-            }
-
-            const value = applyModifiers({
-              value: wrappedOutput,
-              nullable: isNullable,
-              optional: isOptional,
-              nullish: isNullish,
-              defaultValue: meta.default,
-              description: descriptionToApply,
-            })
-
-            if (hasSelfRef) {
-              return `get "${propName}"() { return ${value} }`
-            }
-            return `"${propName}": ${value}`
+        const entries = node.properties.map((prop) => {
+          const { name: propName, schema } = prop
+
+          const meta = ast.syncSchemaRef(schema)
+
+          const isNullable = meta.nullable
+          const isOptional = schema.optional
+          const isNullish = schema.nullish
+
+          const hasSelfRef = this.options.cyclicSchemas != null && ast.containsCircularRef(schema, { circularSchemas: this.options.cyclicSchemas })
+          // Inside a getter the getter itself defers evaluation, so suppress
+          // z.lazy() wrapping on nested refs by temporarily clearing cyclicSchemas.
+          // Save before clearing: this.options === options (same reference via definePrinter),
+          // so reading options.cyclicSchemas after mutation would return undefined.
+          const savedCyclicSchemas = this.options.cyclicSchemas
+          if (hasSelfRef) this.options.cyclicSchemas = undefined
+          const baseOutput = this.transform(schema) ?? this.transform(ast.createSchema({ type: 'unknown' }))!
+          if (hasSelfRef) this.options.cyclicSchemas = savedCyclicSchemas
+
+          const wrappedOutput = this.options.wrapOutput ? this.options.wrapOutput({ output: baseOutput, schema }) || baseOutput : baseOutput
+
+          // When a property schema is not a ref but the metadata is from a ref (e.g., discriminator
+          // property override), skip applying the description from the ref target to avoid applying
+          // metadata from a replaced schema.
+          const descriptionToApply = schema.type !== 'ref' && meta.type === 'ref' ? undefined : meta.description
+
+          const value = applyModifiers({
+            value: wrappedOutput,
+            nullable: isNullable,
+            optional: isOptional,
+            nullish: isNullish,
+            defaultValue: meta.default,
+            description: descriptionToApply,
           })
-          .join(',\n    ')
 
-        let result = `z.object({\n    ${properties}\n    })`
+          if (hasSelfRef) {
+            return `get ${objectKey(propName)}() { return ${value} }`
+          }
+          return `${objectKey(propName)}: ${value}`
+        })
+
+        const objectBase = `z.object(${buildObject(entries)})`
 
         // Handle additionalProperties as .catchall() or .strict()
-        if (node.additionalProperties && node.additionalProperties !== true) {
-          const catchallType = this.transform(node.additionalProperties)
-          if (catchallType) {
-            result += `.catchall(${catchallType})`
+        const result = (() => {
+          if (node.additionalProperties && node.additionalProperties !== true) {
+            const catchallType = this.transform(node.additionalProperties)
+            return catchallType ? `${objectBase}.catchall(${catchallType})` : objectBase
           }
-        } else if (node.additionalProperties === true) {
-          result += `.catchall(${this.transform(ast.createSchema({ type: 'unknown' }))})`
-        } else if (node.additionalProperties === false) {
-          result += '.strict()'
-        }
+          if (node.additionalProperties === true) return `${objectBase}.catchall(${this.transform(ast.createSchema({ type: 'unknown' }))})`
+          if (node.additionalProperties === false) return `${objectBase}.strict()`
+          return objectBase
+        })()
 
         return result
       },
       array(node) {
         const items = (node.items ?? []).map((item) => this.transform(item)).filter(Boolean)
         const inner = items.join(', ') || this.transform(ast.createSchema({ type: 'unknown' }))!
-        let result = `z.array(${inner})${lengthConstraints(node)}`
+        const base = `z.array(${inner})${lengthConstraints(node)}`
 
-        if (node.unique) {
-          result += `.refine(items => new Set(items).size === items.length, { message: "Array entries must be unique" })`
-        }
-
-        return result
+        return node.unique ? `${base}.refine(items => new Set(items).size === items.length, { message: "Array entries must be unique" })` : base
       },
       tuple(node) {
         const items = (node.items ?? []).map((item) => this.transform(item)).filter(Boolean)
 
-        return `z.tuple([${items.join(', ')}])`
+        return `z.tuple(${buildList(items)})`
       },
       union(node) {
         const nodeMembers = node.members ?? []
-        const members = nodeMembers.map((m) => this.transform(m)).filter(Boolean)
+        const members = nodeMembers
+          .map((memberNode) => {
+            const member = this.transform(memberNode)
+
+            return member && node.strategy === 'one' ? strictOneOfMember(member, memberNode) : member
+          })
+          .filter(Boolean)
         if (members.length === 0) return ''
         if (members.length === 1) return members[0]!
         if (node.discriminatorPropertyName && !nodeMembers.some((m) => m.type === 'intersection')) {
           // z.discriminatedUnion requires ZodObject members; intersections (ZodIntersection) are not
           // assignable to $ZodDiscriminant, so fall back to z.union when any member is an intersection.
-          return `z.discriminatedUnion(${stringify(node.discriminatorPropertyName)}, [${members.join(', ')}])`
+          return `z.discriminatedUnion(${stringify(node.discriminatorPropertyName)}, ${buildList(members)})`
         }
 
-        return `z.union([${members.join(', ')}])`
+        return `z.union(${buildList(members)})`
       },
       intersection(node) {
         const members = node.members ?? []
@@ -270,61 +315,37 @@ export const printerZod = ast.definePrinter<PrinterZodFactory>((options) => {
         const [first, ...rest] = members
         if (!first) return ''
 
-        let base = this.transform(first)
-        if (!base) return ''
-
-        for (const member of rest) {
-          if (member.primitive === 'string') {
-            const s = ast.narrowSchema(member, 'string')
-            const c = lengthConstraints(s ?? {})
-            if (c) {
-              base += c
-              continue
-            }
-          } else if (member.primitive === 'number' || member.primitive === 'integer') {
-            const n = ast.narrowSchema(member, 'number') ?? ast.narrowSchema(member, 'integer')
-            const c = numberConstraints(n ?? {})
-            if (c) {
-              base += c
-              continue
-            }
-          } else if (member.primitive === 'array') {
-            const a = ast.narrowSchema(member, 'array')
-            const c = lengthConstraints(a ?? {})
-            if (c) {
-              base += c
-              continue
-            }
-          }
-          const transformed = this.transform(member)
-          if (transformed) base = `${base}.and(${transformed})`
-        }
+        const firstBase = this.transform(first)
+        if (!firstBase) return ''
 
-        return base
+        return rest.reduce((acc, member) => {
+          const constraint = getMemberConstraint(member)
+          if (constraint) return acc + constraint
+          const transformed = this.transform(member)
+          return transformed ? `${acc}.and(${transformed})` : acc
+        }, firstBase)
       },
       ...options.nodes,
     },
     print(node) {
       const { keysToOmit } = this.options
 
-      let base = this.transform(node)
-      if (!base) return null
+      const transformed = this.transform(node)
+      if (!transformed) return null
 
       const meta = ast.syncSchemaRef(node)
 
-      if (keysToOmit?.length && meta.primitive === 'object' && !(meta.type === 'union' && meta.discriminatorPropertyName)) {
+      const base = (() => {
+        if (!keysToOmit?.length || meta.primitive !== 'object' || (meta.type === 'union' && meta.discriminatorPropertyName)) return transformed
         // Mirror printerTs `nonNullable: true`: when omitting keys, the resulting
         // schema is a new non-nullable object type — skip optional/nullable/nullish.
         // Discriminated unions (z.discriminatedUnion) do not support .omit(), so skip them.
 
         // If this is a lazy reference, apply omit inside the lazy function
-        const lazyMatch = base.match(/^z\.lazy\(\(\)\s*=>\s*(.+)\)$/)
-        if (lazyMatch) {
-          base = `z.lazy(() => ${lazyMatch[1]}.omit({ ${keysToOmit.map((k: string) => `"${k}": true`).join(', ')} }))`
-        } else {
-          base = `${base}.omit({ ${keysToOmit.map((k: string) => `"${k}": true`).join(', ')} })`
-        }
-      }
+        const lazyMatch = transformed.match(/^z\.lazy\(\(\)\s*=>\s*(.+)\)$/)
+        if (lazyMatch) return `z.lazy(() => ${lazyMatch[1]}.omit({ ${keysToOmit.map((k: string) => `"${k}": true`).join(', ')} }))`
+        return `${transformed}.omit({ ${keysToOmit.map((k: string) => `"${k}": true`).join(', ')} })`
+      })()
 
       return applyModifiers({
         value: base,

package/src/printers/printerZodMini.ts

@@ -1,4 +1,4 @@
-import { stringify } from '@internals/utils'
+import { buildList, buildObject, extractRefName, objectKey, stringify } from '@kubb/ast/utils'
 
 import { ast } from '@kubb/core'
 import type { PluginZod, ResolverZod } from '../types.ts'
@@ -45,7 +45,7 @@ export type PrinterZodMiniOptions = {
   /**
    * Properties to exclude using `.omit({ key: true })`.
    */
-  keysToOmit?: Array<string>
+  keysToOmit?: Array<string> | null
   /**
    * Schema names that form circular dependency chains.
    * Properties referencing these emit lazy getters wrapping refs in `z.lazy(() => …)`.
@@ -61,6 +61,22 @@ export type PrinterZodMiniOptions = {
  * Factory options for the Zod Mini printer, defining input/output types and configuration.
  */
 export type PrinterZodMiniFactory = ast.PrinterFactoryOptions<'zod-mini', PrinterZodMiniOptions, string, string>
+
+function strictOneOfMember(member: string, node: ast.SchemaNode): string {
+  if (node.type === 'object' && (node.additionalProperties === undefined || node.additionalProperties === false)) {
+    return member.replace(/^z\.object\(/, 'z.strictObject(')
+  }
+
+  return member
+}
+
+function getMemberConstraintMini(member: ast.SchemaNode): string | undefined {
+  if (member.primitive === 'string') return lengthChecksMini(ast.narrowSchema(member, 'string') ?? {}) || undefined
+  if (member.primitive === 'number' || member.primitive === 'integer')
+    return numberChecksMini(ast.narrowSchema(member, 'number') ?? ast.narrowSchema(member, 'integer') ?? {}) || undefined
+  if (member.primitive === 'array') return lengthChecksMini(ast.narrowSchema(member, 'array') ?? {}) || undefined
+}
+
 /**
  * Zod v4 **Mini** printer built with `definePrinter`.
  *
@@ -144,8 +160,8 @@ export const printerZodMini = ast.definePrinter<PrinterZodMiniFactory>((options)
       },
 
       ref(node) {
-        if (!node.name) return undefined
-        const refName = node.ref ? (ast.extractRefName(node.ref) ?? node.name) : node.name
+        if (!node.name) return null
+        const refName = node.ref ? (extractRefName(node.ref) ?? node.name) : node.name
         const resolvedName = node.ref ? (this.options.resolver?.default(refName, 'function') ?? refName) : node.name
 
         if (node.ref && this.options.cyclicSchemas?.has(refName)) {
@@ -155,70 +171,73 @@ export const printerZodMini = ast.definePrinter<PrinterZodMiniFactory>((options)
         return resolvedName
       },
       object(node) {
-        const properties = node.properties
-          .map((prop) => {
-            const { name: propName, schema } = prop
-
-            const meta = ast.syncSchemaRef(schema)
-
-            const isNullable = meta.nullable
-            const isOptional = schema.optional
-            const isNullish = schema.nullish
-
-            const hasSelfRef = this.options.cyclicSchemas != null && ast.containsCircularRef(schema, { circularSchemas: this.options.cyclicSchemas })
-            // Inside a getter the getter itself defers evaluation, so suppress
-            // z.lazy() wrapping on nested refs by temporarily clearing cyclicSchemas.
-            if (hasSelfRef) this.options.cyclicSchemas = undefined
-            const baseOutput = this.transform(schema) ?? this.transform(ast.createSchema({ type: 'unknown' }))!
-            if (hasSelfRef) this.options.cyclicSchemas = options.cyclicSchemas
-
-            const wrappedOutput = this.options.wrapOutput ? this.options.wrapOutput({ output: baseOutput, schema }) || baseOutput : baseOutput
-
-            const value = applyMiniModifiers({
-              value: wrappedOutput,
-              nullable: isNullable,
-              optional: isOptional,
-              nullish: isNullish,
-              defaultValue: meta.default,
-            })
-
-            if (hasSelfRef) {
-              return `get "${propName}"() { return ${value} }`
-            }
-            return `"${propName}": ${value}`
+        const entries = node.properties.map((prop) => {
+          const { name: propName, schema } = prop
+
+          const meta = ast.syncSchemaRef(schema)
+
+          const isNullable = meta.nullable
+          const isOptional = schema.optional
+          const isNullish = schema.nullish
+
+          const hasSelfRef = this.options.cyclicSchemas != null && ast.containsCircularRef(schema, { circularSchemas: this.options.cyclicSchemas })
+          // Inside a getter the getter itself defers evaluation, so suppress
+          // z.lazy() wrapping on nested refs by temporarily clearing cyclicSchemas.
+          // Save before clearing: this.options === options (same reference via definePrinter),
+          // so reading options.cyclicSchemas after mutation would return undefined.
+          const savedCyclicSchemas = this.options.cyclicSchemas
+          if (hasSelfRef) this.options.cyclicSchemas = undefined
+          const baseOutput = this.transform(schema) ?? this.transform(ast.createSchema({ type: 'unknown' }))!
+          if (hasSelfRef) this.options.cyclicSchemas = savedCyclicSchemas
+
+          const wrappedOutput = this.options.wrapOutput ? this.options.wrapOutput({ output: baseOutput, schema }) || baseOutput : baseOutput
+
+          const value = applyMiniModifiers({
+            value: wrappedOutput,
+            nullable: isNullable,
+            optional: isOptional,
+            nullish: isNullish,
+            defaultValue: meta.default,
           })
-          .join(',\n    ')
 
-        return `z.object({\n    ${properties}\n    })`
+          if (hasSelfRef) {
+            return `get ${objectKey(propName)}() { return ${value} }`
+          }
+          return `${objectKey(propName)}: ${value}`
+        })
+
+        return `z.object(${buildObject(entries)})`
       },
       array(node) {
         const items = (node.items ?? []).map((item) => this.transform(item)).filter(Boolean)
         const inner = items.join(', ') || this.transform(ast.createSchema({ type: 'unknown' }))!
-        let result = `z.array(${inner})${lengthChecksMini(node)}`
+        const base = `z.array(${inner})${lengthChecksMini(node)}`
 
-        if (node.unique) {
-          result += `.refine(items => new Set(items).size === items.length, { message: "Array entries must be unique" })`
-        }
-
-        return result
+        return node.unique ? `${base}.refine(items => new Set(items).size === items.length, { message: "Array entries must be unique" })` : base
       },
       tuple(node) {
         const items = (node.items ?? []).map((item) => this.transform(item)).filter(Boolean)
 
-        return `z.tuple([${items.join(', ')}])`
+        return `z.tuple(${buildList(items)})`
       },
       union(node) {
         const nodeMembers = node.members ?? []
-        const members = nodeMembers.map((m) => this.transform(m)).filter(Boolean)
+        const members = nodeMembers
+          .map((memberNode) => {
+            const member = this.transform(memberNode)
+
+            return member && node.strategy === 'one' ? strictOneOfMember(member, memberNode) : member
+          })
+          .filter(Boolean)
         if (members.length === 0) return ''
         if (members.length === 1) return members[0]!
         if (node.discriminatorPropertyName && !nodeMembers.some((m) => m.type === 'intersection')) {
           // z.discriminatedUnion requires ZodObject members; intersections (ZodIntersection) are not
           // assignable to $ZodDiscriminant, so fall back to z.union when any member is an intersection.
-          return `z.discriminatedUnion(${stringify(node.discriminatorPropertyName)}, [${members.join(', ')}])`
+          return `z.discriminatedUnion(${stringify(node.discriminatorPropertyName)}, ${buildList(members)})`
         }
 
-        return `z.union([${members.join(', ')}])`
+        return `z.union(${buildList(members)})`
       },
       intersection(node) {
         const members = node.members ?? []
@@ -227,61 +246,37 @@ export const printerZodMini = ast.definePrinter<PrinterZodMiniFactory>((options)
         const [first, ...rest] = members
         if (!first) return ''
 
-        let base = this.transform(first)
-        if (!base) return ''
-
-        for (const member of rest) {
-          if (member.primitive === 'string') {
-            const s = ast.narrowSchema(member, 'string')
-            const c = lengthChecksMini(s ?? {})
-            if (c) {
-              base += c
-              continue
-            }
-          } else if (member.primitive === 'number' || member.primitive === 'integer') {
-            const n = ast.narrowSchema(member, 'number') ?? ast.narrowSchema(member, 'integer')
-            const c = numberChecksMini(n ?? {})
-            if (c) {
-              base += c
-              continue
-            }
-          } else if (member.primitive === 'array') {
-            const a = ast.narrowSchema(member, 'array')
-            const c = lengthChecksMini(a ?? {})
-            if (c) {
-              base += c
-              continue
-            }
-          }
-          const transformed = this.transform(member)
-          if (transformed) base = `z.intersection(${base}, ${transformed})`
-        }
+        const firstBase = this.transform(first)
+        if (!firstBase) return ''
 
-        return base
+        return rest.reduce((acc, member) => {
+          const constraint = getMemberConstraintMini(member)
+          if (constraint) return acc + constraint
+          const transformed = this.transform(member)
+          return transformed ? `z.intersection(${acc}, ${transformed})` : acc
+        }, firstBase)
       },
       ...options.nodes,
     },
     print(node) {
       const { keysToOmit } = this.options
 
-      let base = this.transform(node)
-      if (!base) return null
+      const transformed = this.transform(node)
+      if (!transformed) return null
 
       const meta = ast.syncSchemaRef(node)
 
-      if (keysToOmit?.length && meta.primitive === 'object' && !(meta.type === 'union' && meta.discriminatorPropertyName)) {
+      const base = (() => {
+        if (!keysToOmit?.length || meta.primitive !== 'object' || (meta.type === 'union' && meta.discriminatorPropertyName)) return transformed
         // Mirror printerTs `nonNullable: true`: when omitting keys, the resulting
         // schema is a new non-nullable object type — skip optional/nullable/nullish.
         // Discriminated unions (z.discriminatedUnion) do not support .omit(), so skip them.
 
         // If this is a lazy reference, apply omit inside the lazy function
-        const lazyMatch = base.match(/^z\.lazy\(\(\)\s*=>\s*(.+)\)$/)
-        if (lazyMatch) {
-          base = `z.lazy(() => ${lazyMatch[1]}.omit({ ${keysToOmit.map((k: string) => `"${k}": true`).join(', ')} }))`
-        } else {
-          base = `${base}.omit({ ${keysToOmit.map((k: string) => `"${k}": true`).join(', ')} })`
-        }
-      }
+        const lazyMatch = transformed.match(/^z\.lazy\(\(\)\s*=>\s*(.+)\)$/)
+        if (lazyMatch) return `z.lazy(() => ${lazyMatch[1]}.omit({ ${keysToOmit.map((k: string) => `"${k}": true`).join(', ')} }))`
+        return `${transformed}.omit({ ${keysToOmit.map((k: string) => `"${k}": true`).join(', ')} })`
+      })()
 
       return applyMiniModifiers({
         value: base,