feat: add support for discriminatedUnion (#281)

* fix: remove duplicate test

* feat: add @Discriminator tag

* chore: remove codecov

* doc: update

README.md

@@ -49,15 +49,15 @@ By default, `FormatType` is defined as the following type (corresponding Zod val
 ```ts
 type FormatType =
   | "date-time" // z.string().datetime()
-  | "date"      // z.string().date()
-  | "time"      // z.string().time()
-  | "duration"  // z.string().duration()
-  | "email"     // z.string().email()
-  | "ip"        // z.string().ip()
-  | "ipv4"      // z.string().ip()
-  | "ipv6"      // z.string().ip()
-  | "url"       // z.string().url()
-  | "uuid";     // z.string().uuid()
+  | "date" // z.string().date()
+  | "time" // z.string().time()
+  | "duration" // z.string().duration()
+  | "email" // z.string().email()
+  | "ip" // z.string().ip()
+  | "ipv4" // z.string().ip()
+  | "ipv6" // z.string().ip()
+  | "url" // z.string().url()
+  | "uuid"; // z.string().uuid()
 ```
 
 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({
 Other JSDoc tags are available:
 
 | JSDoc keyword          | JSDoc Example            | Description                                                                                                                               | Generated Zod                      |
-|------------------------|--------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------|
+| ---------------------- | ------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------- |
 | `@description {value}` | `@description Full name` | Sets the description of the property                                                                                                      | `z.string().describe("Full name")` |
 | `@default {value}`     | `@default 42`            | Sets a default value for the property                                                                                                     | `z.number().default(42)`           |
 | `@strict`              | `@strict`                | Adds the `strict()` modifier to an object                                                                                                 | `z.object().strict()`              |
 | `@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')`          |
-  |
+
+## JSDoc tag for `union` types
+
+| JSDoc keyword               | JSDoc Example       | Description                                           | Generated Zod                     |
+| --------------------------- | ------------------- | ----------------------------------------------------- | --------------------------------- |
+| `@discriminator {propName}` | `@discriminator id` | Generates a `z.discriminatedUnion()` instead of union | `z.discriminatedUnion("id", ...)` |
+
+Example:
+
+```ts
+// source.ts
+/**
+ * @discriminator type
+ **/
+export type Person =
+  | { type: "Adult"; name: string }
+  | { type: "Child"; age: number };
+
+// output.ts
+// Generated by ts-to-zod
+import { z } from "zod";
+
+export const personSchema = z.discriminatedUnion("type", [
+  z.object({
+    type: z.literal("Adult"),
+    name: z.string(),
+  }),
+  z.object({
+    type: z.literal("Child"),
+    age: z.number(),
+  }),
+]);
+```
 
 ## JSDoc tags for elements of `string` and `number` arrays
 

src/core/generateZodSchema.test.ts

@@ -1492,6 +1492,83 @@ describe("generateZodSchema", () => {
     `);
   });
 
+  it("should generate a discriminatedUnion when @discriminator is used", () => {
+    const source = `
+    /**
+     * @discriminator id
+     **/
+    export type A = { id: "1"; name: string; } | { id: "2"; age: number; }`;
+
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "/**
+       * @discriminator id
+       **/
+      export const aSchema = z.discriminatedUnion("id", [z.object({
+              id: z.literal("1"),
+              name: z.string()
+          }), z.object({
+              id: z.literal("2"),
+              age: z.number()
+          })]);"
+    `);
+  });
+
+  it("should generate a discriminatedUnion with a referenced type", () => {
+    const source = `
+    /**
+     * @discriminator id
+     **/
+    export type Foo = { id: "1"; name: string; } | Bar`;
+
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "/**
+       * @discriminator id
+       **/
+      export const fooSchema = z.discriminatedUnion("id", [z.object({
+              id: z.literal("1"),
+              name: z.string()
+          }), barSchema]);"
+    `);
+  });
+
+  it("should fall back to union when types are not discriminated", () => {
+    const source = `
+    /**
+     * @discriminator id
+     **/
+    export type A = { id: "1"; name: string; } | string`;
+
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "/**
+       * @discriminator id
+       **/
+      export const aSchema = z.union([z.object({
+              id: z.literal("1"),
+              name: z.string()
+          }), z.string()]);"
+    `);
+  });
+
+  it("should fall back to union when discriminator is missing", () => {
+    const source = `
+    /**
+     * @discriminator id
+     **/
+    export type A = { name: string; } | { id: "2"; age: number; }`;
+
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "/**
+       * @discriminator id
+       **/
+      export const aSchema = z.union([z.object({
+              name: z.string()
+          }), z.object({
+              id: z.literal("2"),
+              age: z.number()
+          })]);"
+    `);
+  });
+
   it("should deal with @default with all types", () => {
     const source = `export interface WithDefaults {
      /**
@@ -1681,15 +1758,6 @@ describe("generateZodSchema", () => {
     `);
   });
 
-  it("should throw on generics", () => {
-    const source = `export interface Villain<TPower> {
-     powers: TPower[]
-   }`;
-    expect(() => generate(source)).toThrowErrorMatchingInlineSnapshot(
-      `"Interface with generics are not supported!"`
-    );
-  });
-
   it("should throw on interface with generics", () => {
     const source = `export interface Villain<TPower> {
      powers: TPower[]

src/core/generateZodSchema.ts

@@ -650,6 +650,56 @@ function buildZodPrimitiveInternal({
       });
     }
 
+    if (jsDocTags.discriminator) {
+      let isValidDiscriminatedUnion = true;
+
+      // Check each member of the union
+      for (const node of nodes) {
+        if (!ts.isTypeLiteralNode(node) && !ts.isTypeReferenceNode(node)) {
+          console.warn(
+            ` »   Warning: discriminated union member "${node.getText(
+              sourceFile
+            )}" is not a type reference or object literal`
+          );
+          isValidDiscriminatedUnion = false;
+          break;
+        }
+
+        // For type references, we'd need to resolve the referenced type
+        // For type literals, we can check directly
+        if (ts.isTypeLiteralNode(node)) {
+          const hasDiscriminator = node.members.some(
+            (member) =>
+              ts.isPropertySignature(member) &&
+              member.name &&
+              member.name.getText(sourceFile) === jsDocTags.discriminator
+          );
+
+          if (!hasDiscriminator) {
+            console.warn(
+              ` »   Warning: discriminated union member "${node.getText(
+                sourceFile
+              )}" missing discriminator field "${jsDocTags.discriminator}"`
+            );
+            isValidDiscriminatedUnion = false;
+            break;
+          }
+        }
+      }
+
+      if (isValidDiscriminatedUnion) {
+        return buildZodSchema(
+          z,
+          "discriminatedUnion",
+          [
+            f.createStringLiteral(jsDocTags.discriminator),
+            f.createArrayLiteralExpression(values),
+          ],
+          zodProperties
+        );
+      }
+    }
+
     return buildZodSchema(
       z,
       "union",

src/core/jsDocTags.ts

@@ -74,6 +74,7 @@ export interface JSDocTagsBase {
   pattern?: string;
   strict?: boolean;
   schema?: string;
+  discriminator?: string;
 }
 
 export type ElementJSDocTags = Pick<
@@ -108,6 +109,7 @@ const jsDocTagKeys: Array<keyof JSDocTags> = [
   "elementMaxLength",
   "elementPattern",
   "elementFormat",
+  "discriminator",
 ];
 
 /**
@@ -207,6 +209,9 @@ export function getJSDocTags(nodeType: ts.Node, sourceFile: ts.SourceFile) {
               }
             }
             break;
+          case "discriminator":
+            jsDocTags[tagName] = tag.comment;
+            break;
           case "strict":
             break;
           default: