refactor: generateExample with zod union handler

Author: omermecitoglu

src/utils/generateExample.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, it } from "@jest/globals";
-import z from "zod";
+import z, { type ZodType } from "zod";
 import generateExample from "./generateExample";
 
 describe("generateExample", () => {
@@ -104,6 +104,14 @@ describe("generateExample", () => {
   it("should generate an example from a zod array", () => {
     expect(generateExample(z.string().array(), true)).toStrictEqual(["example string"]);
     expect(generateExample(z.number().array(), true)).toStrictEqual([1]);
+    expect(generateExample(z.array(z.union([z.string(), z.number()])), true)).toStrictEqual(["example string", 1]);
+  });
+
+  it("should generate an example from a zod union", () => {
+    const schema = z.union([z.string(), z.number()]);
+    expect(generateExample(schema, true)).toBe("example string");
+    const emptyUnion = z.union([] as unknown as [ZodType, ZodType]);
+    expect(generateExample(emptyUnion, true)).toBe(undefined);
   });
 
   it("should throw an error from an unknown zod schema", () => {

src/utils/generateExample.ts

@@ -1,4 +1,4 @@
-import { ZodArray, ZodBigInt, ZodBoolean, ZodEnum, ZodNull, ZodNullable, ZodNumber, ZodObject, ZodOptional, ZodString, type ZodType, type ZodTypeDef, ZodUndefined } from "zod";
+import { ZodArray, ZodBigInt, ZodBoolean, ZodEnum, ZodNull, ZodNullable, ZodNumber, ZodObject, ZodOptional, ZodString, type ZodType, type ZodTypeDef, ZodUndefined, ZodUnion } from "zod";
 
 export default function generateExample<I, O>(
   schema: ZodType<O, ZodTypeDef, I>,
@@ -100,7 +100,19 @@ export default function generateExample<I, O>(
     return "example string" as unknown as O;
   }
   if (schema instanceof ZodArray) {
+    if (schema._def.type instanceof ZodUnion) {
+      const options = schema._def.type.options as ZodType[];
+      return options.map(option => generateExample(option, ignoreOptionals)) as unknown as O;
+    }
     return [generateExample(schema._def.type, ignoreOptionals)] as unknown as O;
   }
+  if (schema instanceof ZodUnion) {
+    const options = schema.options as ZodType[];
+    const firstItem = options[0];
+    if (firstItem) {
+      return generateExample(firstItem, ignoreOptionals);
+    }
+    return undefined as unknown as O;
+  }
   throw new Error("Unknown zod schema");
 }