Skip to content

Commit 4dfbab5

Browse files
tvillarentvillaren
authored
feat: add support for discriminatedUnion (#281)
* fix: remove duplicate test * feat: add @Discriminator tag * chore: remove codecov * doc: update
1 parent 56b4374 commit 4dfbab5

File tree

4 files changed

+175
-20
lines changed

4 files changed

+175
-20
lines changed

‎README.md

+43-11
Original file line numberDiff line numberDiff line change
@@ -49,15 +49,15 @@ By default, `FormatType` is defined as the following type (corresponding Zod val
4949
```ts
5050
type FormatType =
5151
| "date-time" // z.string().datetime()
52-
| "date" // z.string().date()
53-
| "time" // z.string().time()
54-
| "duration" // z.string().duration()
55-
| "email" // z.string().email()
56-
| "ip" // z.string().ip()
57-
| "ipv4" // z.string().ip()
58-
| "ipv6" // z.string().ip()
59-
| "url" // z.string().url()
60-
| "uuid"; // z.string().uuid()
52+
| "date" // z.string().date()
53+
| "time" // z.string().time()
54+
| "duration" // z.string().duration()
55+
| "email" // z.string().email()
56+
| "ip" // z.string().ip()
57+
| "ipv4" // z.string().ip()
58+
| "ipv6" // z.string().ip()
59+
| "url" // z.string().url()
60+
| "uuid"; // z.string().uuid()
6161
```
6262

6363
However, see the section on [Custom JSDoc Format Types](#custom-jsdoc-format-types) to learn more about defining other types of formats for string validation.
@@ -151,12 +151,44 @@ export const heroContactSchema = z.object({
151151
Other JSDoc tags are available:
152152

153153
| JSDoc keyword | JSDoc Example | Description | Generated Zod |
154-
|------------------------|--------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------|
154+
| ---------------------- | ------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------- |
155155
| `@description {value}` | `@description Full name` | Sets the description of the property | `z.string().describe("Full name")` |
156156
| `@default {value}` | `@default 42` | Sets a default value for the property | `z.number().default(42)` |
157157
| `@strict` | `@strict` | Adds the `strict()` modifier to an object | `z.object().strict()` |
158158
| `@schema` | `@schema .catch('foo')` | If value starts with a `.`, appends the specified value to the generated schema. Otherwise this value will override the generated schema. | `z.string().catch('foo')` |
159-
|
159+
160+
## JSDoc tag for `union` types
161+
162+
| JSDoc keyword | JSDoc Example | Description | Generated Zod |
163+
| --------------------------- | ------------------- | ----------------------------------------------------- | --------------------------------- |
164+
| `@discriminator {propName}` | `@discriminator id` | Generates a `z.discriminatedUnion()` instead of union | `z.discriminatedUnion("id", ...)` |
165+
166+
Example:
167+
168+
```ts
169+
// source.ts
170+
/**
171+
* @discriminator type
172+
**/
173+
export type Person =
174+
| { type: "Adult"; name: string }
175+
| { type: "Child"; age: number };
176+
177+
// output.ts
178+
// Generated by ts-to-zod
179+
import { z } from "zod";
180+
181+
export const personSchema = z.discriminatedUnion("type", [
182+
z.object({
183+
type: z.literal("Adult"),
184+
name: z.string(),
185+
}),
186+
z.object({
187+
type: z.literal("Child"),
188+
age: z.number(),
189+
}),
190+
]);
191+
```
160192

161193
## JSDoc tags for elements of `string` and `number` arrays
162194

‎src/core/generateZodSchema.test.ts

+77-9
Original file line numberDiff line numberDiff line change
@@ -1492,6 +1492,83 @@ describe("generateZodSchema", () => {
14921492
`);
14931493
});
14941494

1495+
it("should generate a discriminatedUnion when @discriminator is used", () => {
1496+
const source = `
1497+
/**
1498+
* @discriminator id
1499+
**/
1500+
export type A = { id: "1"; name: string; } | { id: "2"; age: number; }`;
1501+
1502+
expect(generate(source)).toMatchInlineSnapshot(`
1503+
"/**
1504+
* @discriminator id
1505+
**/
1506+
export const aSchema = z.discriminatedUnion("id", [z.object({
1507+
id: z.literal("1"),
1508+
name: z.string()
1509+
}), z.object({
1510+
id: z.literal("2"),
1511+
age: z.number()
1512+
})]);"
1513+
`);
1514+
});
1515+
1516+
it("should generate a discriminatedUnion with a referenced type", () => {
1517+
const source = `
1518+
/**
1519+
* @discriminator id
1520+
**/
1521+
export type Foo = { id: "1"; name: string; } | Bar`;
1522+
1523+
expect(generate(source)).toMatchInlineSnapshot(`
1524+
"/**
1525+
* @discriminator id
1526+
**/
1527+
export const fooSchema = z.discriminatedUnion("id", [z.object({
1528+
id: z.literal("1"),
1529+
name: z.string()
1530+
}), barSchema]);"
1531+
`);
1532+
});
1533+
1534+
it("should fall back to union when types are not discriminated", () => {
1535+
const source = `
1536+
/**
1537+
* @discriminator id
1538+
**/
1539+
export type A = { id: "1"; name: string; } | string`;
1540+
1541+
expect(generate(source)).toMatchInlineSnapshot(`
1542+
"/**
1543+
* @discriminator id
1544+
**/
1545+
export const aSchema = z.union([z.object({
1546+
id: z.literal("1"),
1547+
name: z.string()
1548+
}), z.string()]);"
1549+
`);
1550+
});
1551+
1552+
it("should fall back to union when discriminator is missing", () => {
1553+
const source = `
1554+
/**
1555+
* @discriminator id
1556+
**/
1557+
export type A = { name: string; } | { id: "2"; age: number; }`;
1558+
1559+
expect(generate(source)).toMatchInlineSnapshot(`
1560+
"/**
1561+
* @discriminator id
1562+
**/
1563+
export const aSchema = z.union([z.object({
1564+
name: z.string()
1565+
}), z.object({
1566+
id: z.literal("2"),
1567+
age: z.number()
1568+
})]);"
1569+
`);
1570+
});
1571+
14951572
it("should deal with @default with all types", () => {
14961573
const source = `export interface WithDefaults {
14971574
/**
@@ -1681,15 +1758,6 @@ describe("generateZodSchema", () => {
16811758
`);
16821759
});
16831760

1684-
it("should throw on generics", () => {
1685-
const source = `export interface Villain<TPower> {
1686-
powers: TPower[]
1687-
}`;
1688-
expect(() => generate(source)).toThrowErrorMatchingInlineSnapshot(
1689-
`"Interface with generics are not supported!"`
1690-
);
1691-
});
1692-
16931761
it("should throw on interface with generics", () => {
16941762
const source = `export interface Villain<TPower> {
16951763
powers: TPower[]

‎src/core/generateZodSchema.ts

+50
Original file line numberDiff line numberDiff line change
@@ -650,6 +650,56 @@ function buildZodPrimitiveInternal({
650650
});
651651
}
652652

653+
if (jsDocTags.discriminator) {
654+
let isValidDiscriminatedUnion = true;
655+
656+
// Check each member of the union
657+
for (const node of nodes) {
658+
if (!ts.isTypeLiteralNode(node) && !ts.isTypeReferenceNode(node)) {
659+
console.warn(
660+
` » Warning: discriminated union member "${node.getText(
661+
sourceFile
662+
)}" is not a type reference or object literal`
663+
);
664+
isValidDiscriminatedUnion = false;
665+
break;
666+
}
667+
668+
// For type references, we'd need to resolve the referenced type
669+
// For type literals, we can check directly
670+
if (ts.isTypeLiteralNode(node)) {
671+
const hasDiscriminator = node.members.some(
672+
(member) =>
673+
ts.isPropertySignature(member) &&
674+
member.name &&
675+
member.name.getText(sourceFile) === jsDocTags.discriminator
676+
);
677+
678+
if (!hasDiscriminator) {
679+
console.warn(
680+
` » Warning: discriminated union member "${node.getText(
681+
sourceFile
682+
)}" missing discriminator field "${jsDocTags.discriminator}"`
683+
);
684+
isValidDiscriminatedUnion = false;
685+
break;
686+
}
687+
}
688+
}
689+
690+
if (isValidDiscriminatedUnion) {
691+
return buildZodSchema(
692+
z,
693+
"discriminatedUnion",
694+
[
695+
f.createStringLiteral(jsDocTags.discriminator),
696+
f.createArrayLiteralExpression(values),
697+
],
698+
zodProperties
699+
);
700+
}
701+
}
702+
653703
return buildZodSchema(
654704
z,
655705
"union",

‎src/core/jsDocTags.ts

+5
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ export interface JSDocTagsBase {
7474
pattern?: string;
7575
strict?: boolean;
7676
schema?: string;
77+
discriminator?: string;
7778
}
7879

7980
export type ElementJSDocTags = Pick<
@@ -108,6 +109,7 @@ const jsDocTagKeys: Array<keyof JSDocTags> = [
108109
"elementMaxLength",
109110
"elementPattern",
110111
"elementFormat",
112+
"discriminator",
111113
];
112114

113115
/**
@@ -207,6 +209,9 @@ export function getJSDocTags(nodeType: ts.Node, sourceFile: ts.SourceFile) {
207209
}
208210
}
209211
break;
212+
case "discriminator":
213+
jsDocTags[tagName] = tag.comment;
214+
break;
210215
case "strict":
211216
break;
212217
default:

0 commit comments

Comments
 (0)