Skip to content

Commit 0d5cb25

Browse files
franklevasseurfranklevasseur
authored
feat(zui): better support for discriminated unions transformation (#664)
1 parent a086be8 commit 0d5cb25

File tree

8 files changed

+89
-22
lines changed

8 files changed

+89
-22
lines changed

‎zui/package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "@bpinternal/zui",
3-
"version": "1.2.3",
3+
"version": "1.3.0",
44
"description": "A fork of Zod with additional features",
55
"type": "module",
66
"source": "./src/index.ts",

‎zui/src/transforms/common/json-schema.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,10 @@ type NullableDef = util.Satisfies<{ typeName: z.ZodFirstPartyTypeKind.ZodNullabl
1313
type OptionalDef = util.Satisfies<{ typeName: z.ZodFirstPartyTypeKind.ZodOptional }, Partial<z.ZodOptionalDef>>
1414
type UndefinedDef = util.Satisfies<{ typeName: z.ZodFirstPartyTypeKind.ZodUndefined }, Partial<z.ZodUndefinedDef>>
1515
type UnknownDef = util.Satisfies<{ typeName: z.ZodFirstPartyTypeKind.ZodUnknown }, Partial<z.ZodUnknownDef>>
16+
type DiscriminatedUnionDef = util.Satisfies<
17+
{ typeName: z.ZodFirstPartyTypeKind.ZodDiscriminatedUnion; discriminator?: string },
18+
Partial<z.ZodDiscriminatedUnionDef>
19+
>
1620

1721
/**
1822
* ZuiJSONSchema:
@@ -99,7 +103,7 @@ export type AnySchema = BaseZuiJSONSchema
99103
export type UnknownSchema = BaseZuiJSONSchema<UnknownDef>
100104
export type ArraySchema = _ArraySchema & BaseZuiJSONSchema
101105
export type UnionSchema = _UnionSchema & BaseZuiJSONSchema
102-
export type DiscriminatedUnionSchema = _DiscriminatedUnionSchema & BaseZuiJSONSchema
106+
export type DiscriminatedUnionSchema = _DiscriminatedUnionSchema & BaseZuiJSONSchema<DiscriminatedUnionDef>
103107
export type IntersectionSchema = _IntersectionSchema & BaseZuiJSONSchema
104108
export type SetSchema = _SetSchema & BaseZuiJSONSchema
105109
export type EnumSchema = _EnumSchema & BaseZuiJSONSchema

‎zui/src/transforms/transform-pipeline.test.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,13 @@ describe.concurrent('transformPipeline', () => {
379379
const srcSchema = z.union([z.string(), z.number()])
380380
assert(srcSchema).toTransformBackToItself()
381381
})
382+
it('should map ZodDiscriminatedUnion to itself', async () => {
383+
const srcSchema = z.discriminatedUnion('type', [
384+
z.object({ type: z.literal('foo'), foo: z.string() }).strict(),
385+
z.object({ type: z.literal('bar'), bar: z.number() }).strict(),
386+
])
387+
assert(srcSchema).toTransformBackToItself()
388+
})
382389
it('should map ZodIntersection to itself', async () => {
383390
const srcSchema = z.intersection(
384391
z.object({ type: z.literal('foo'), foo: z.string() }),

‎zui/src/transforms/zui-from-json-schema/guards.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,3 +19,7 @@ export const isUndefinedSchema = (s: JSONSchema7): s is json.UndefinedSchema =>
1919

2020
export const isUnknownSchema = (s: JSONSchema7): s is json.UnknownSchema =>
2121
!s.not && (s as json.UnknownSchema)['x-zui']?.def?.typeName === z.ZodFirstPartyTypeKind.ZodUnknown
22+
23+
export const isDiscriminatedUnionSchema = (s: JSONSchema7): s is json.DiscriminatedUnionSchema =>
24+
s.anyOf !== undefined &&
25+
(s as json.DiscriminatedUnionSchema)['x-zui']?.def?.typeName === z.ZodFirstPartyTypeKind.ZodDiscriminatedUnion

‎zui/src/transforms/zui-from-json-schema/index.test.ts‎

Lines changed: 53 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -185,23 +185,26 @@ describe.concurrent('zuifromJSONSchemaNext', () => {
185185
assert(zSchema).toEqual(expected)
186186
})
187187

188-
test('should map DiscriminatedUnionSchema to ZodUnion', () => {
189-
const jSchema = buildSchema({
190-
anyOf: [
191-
{
192-
type: 'object',
193-
properties: { type: { type: 'string', const: 'A' }, a: { type: 'string' } },
194-
required: ['type', 'a'],
195-
},
196-
{
197-
type: 'object',
198-
properties: { type: { type: 'string', const: 'B' }, b: { type: 'number' } },
199-
required: ['type', 'b'],
200-
},
201-
],
202-
})
188+
test('should map DiscriminatedUnionSchema to ZodDiscriminatedUnion', () => {
189+
const jSchema = buildSchema(
190+
{
191+
anyOf: [
192+
{
193+
type: 'object',
194+
properties: { type: { type: 'string', const: 'A' }, a: { type: 'string' } },
195+
required: ['type', 'a'],
196+
},
197+
{
198+
type: 'object',
199+
properties: { type: { type: 'string', const: 'B' }, b: { type: 'number' } },
200+
required: ['type', 'b'],
201+
},
202+
],
203+
},
204+
{ def: { typeName: 'ZodDiscriminatedUnion', discriminator: 'type' } },
205+
)
203206
const zSchema = fromJSONSchema(jSchema)
204-
const expected = z.union([
207+
const expected = z.discriminatedUnion('type', [
205208
z.object({ type: z.literal('A'), a: z.string() }),
206209
z.object({ type: z.literal('B'), b: z.number() }),
207210
])
@@ -662,7 +665,7 @@ describe.concurrent('zuifromJSONSchemaNext', () => {
662665
expect(getTypescriptType(original)).toBe(getTypescriptType(restored))
663666
})
664667

665-
test('should preserve discriminated union with descriptions', () => {
668+
test('should preserve union with descriptions', () => {
666669
const original = z
667670
.union([
668671
z.object({ type: z.literal('success'), data: z.string().describe('Success data') }),
@@ -676,6 +679,20 @@ describe.concurrent('zuifromJSONSchemaNext', () => {
676679
expect(restored._def.description).toBe('API Response')
677680
})
678681

682+
test('should preserve discriminated union with descriptions', () => {
683+
const original = z
684+
.discriminatedUnion('type', [
685+
z.object({ type: z.literal('success'), data: z.string().describe('Success data') }),
686+
z.object({ type: z.literal('error'), message: z.string().describe('Error message') }),
687+
])
688+
.describe('API Response')
689+
690+
const restored = roundTrip(original)
691+
692+
expect(getTypescriptType(original)).toBe(getTypescriptType(restored))
693+
expect(restored._def.description).toBe('API Response')
694+
})
695+
679696
test('should preserve nullable fields', () => {
680697
const original = z.object({
681698
name: z.string(),
@@ -930,7 +947,7 @@ describe.concurrent('zuifromJSONSchemaNext', () => {
930947
expect(getTypescriptType(original)).toBe(getTypescriptType(restored))
931948
})
932949

933-
test('should preserve complex discriminated union structure', () => {
950+
test('should preserve complex union structure', () => {
934951
const original = z.union([
935952
z.object({
936953
kind: z.literal('circle'),
@@ -951,6 +968,24 @@ describe.concurrent('zuifromJSONSchemaNext', () => {
951968
expect(getTypescriptType(original)).toBe(getTypescriptType(restored))
952969
})
953970

971+
test('should preserve discriminated union of objects with optional fields', () => {
972+
const original = z.discriminatedUnion('type', [
973+
z.object({
974+
type: z.literal('A'),
975+
value: z.string(),
976+
extra: z.number().optional(),
977+
}),
978+
z.object({
979+
type: z.literal('B'),
980+
data: z.boolean(),
981+
meta: z.string().optional(),
982+
}),
983+
])
984+
const restored = roundTrip(original)
985+
986+
expect(getTypescriptType(original)).toBe(getTypescriptType(restored))
987+
})
988+
954989
test('should preserve union of objects with optional fields', () => {
955990
const original = z.union([
956991
z.object({

‎zui/src/transforms/zui-from-json-schema/index.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,7 @@ function _fromJSONSchema(schema: JSONSchema7Definition | undefined): z.ZodType {
156156
if (schema.anyOf.length === 0) {
157157
return DEFAULT_TYPE
158158
}
159+
159160
if (schema.anyOf.length === 1) {
160161
return _fromJSONSchema(schema.anyOf[0])
161162
}
@@ -170,6 +171,16 @@ function _fromJSONSchema(schema: JSONSchema7Definition | undefined): z.ZodType {
170171
return inner.nullable()
171172
}
172173

174+
if (guards.isDiscriminatedUnionSchema(schema) && schema['x-zui']?.def?.discriminator) {
175+
const { discriminator } = schema['x-zui'].def
176+
const options = schema.anyOf.map(_fromJSONSchema) as [
177+
z.ZodDiscriminatedUnionOption<string>,
178+
z.ZodDiscriminatedUnionOption<string>,
179+
...z.ZodDiscriminatedUnionOption<string>[],
180+
]
181+
return z.discriminatedUnion(discriminator, options)
182+
}
183+
173184
const options = schema.anyOf.map(_fromJSONSchema) as [z.ZodType, z.ZodType, ...z.ZodType[]]
174185
return z.union(options)
175186
}

‎zui/src/transforms/zui-to-json-schema/index.test.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,9 @@ describe('zuiToJSONSchemaNext', () => {
191191
additionalProperties: false,
192192
},
193193
],
194+
'x-zui': {
195+
def: { typeName: z.ZodFirstPartyTypeKind.ZodDiscriminatedUnion, discriminator: 'type' },
196+
},
194197
})
195198
})
196199

‎zui/src/transforms/zui-to-json-schema/index.ts‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,8 +109,11 @@ export function toJSONSchema(schema: z.Schema): json.Schema {
109109
return {
110110
description: def.description,
111111
anyOf: def.options.map((option) => toJSONSchema(option)),
112-
'x-zui': def['x-zui'],
113-
} satisfies json.UnionSchema
112+
'x-zui': {
113+
...def['x-zui'],
114+
def: { typeName: z.ZodFirstPartyTypeKind.ZodDiscriminatedUnion, discriminator: def.discriminator },
115+
},
116+
} satisfies json.DiscriminatedUnionSchema
114117

115118
case z.ZodFirstPartyTypeKind.ZodIntersection:
116119
const left = toJSONSchema(def.left)

0 commit comments

Comments
 (0)