Merge pull request #129 from A5rocks/summarize-union-of-literals

Turn unions of literals into enums
anatine · Jun 16, 2023 · eb1756d · eb1756d
2 parents aa7b7d8 + 2b4981b
commit eb1756d
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 3 deletions.
diff --git a/packages/zod-openapi/src/lib/zod-openapi.spec.ts b/packages/zod-openapi/src/lib/zod-openapi.spec.ts
@@ -823,4 +823,35 @@ describe('zodOpenapi', () => {
       format: 'binary',
     });
   });
+
+  it('can summarize unions of zod literals as an enum', () => {
+    expect(generateSchema(z.union([z.literal('h'), z.literal('i')]))).toEqual({
+      type: 'string',
+      enum: ['h', 'i']
+    });
+
+    expect(generateSchema(z.union([z.literal(3), z.literal(4)]))).toEqual({
+      type: 'number',
+      enum: [3, 4]
+    });
+
+    // should this just remove the enum? true | false is exhaustive...
+    expect(generateSchema(z.union([z.literal(true), z.literal(false)]))).toEqual({
+      type: 'boolean',
+      enum: [true, false]
+    });
+
+    expect(generateSchema(z.union([z.literal(5), z.literal('i')]))).toEqual({
+      oneOf: [
+        {
+          type: 'number',
+          enum: [5]
+        },
+        {
+          type: 'string',
+          enum: ['i']
+        }
+      ]
+    });
+  })
 });
diff --git a/packages/zod-openapi/src/lib/zod-openapi.ts b/packages/zod-openapi/src/lib/zod-openapi.ts
@@ -388,11 +388,33 @@ function parseUnion({
   zodRef,
   useOutput,
 }: ParsingArgs<z.ZodUnion<[z.ZodTypeAny, ...z.ZodTypeAny[]]>>): SchemaObject {
+  const contents = zodRef._def.options;
+  if (contents.reduce((prev, content) => prev && content._def.typeName === 'ZodLiteral', true)) {
+    // special case to transform unions of literals into enums
+    const literals = contents as unknown as z.ZodLiteral<OpenApiZodAny>[];
+    const type = literals
+      .reduce((prev, content) =>
+        !prev || prev === typeof content._def.value ?
+          typeof content._def.value :
+          null,
+        null as null | string
+      );
+
+    if (type) {
+      return merge(
+        {
+          type: type as 'string' | 'number' | 'boolean',
+          enum: literals.map((literal) => literal._def.value)
+        },
+        zodRef.description ? { description: zodRef.description } : {},
+        ...schemas
+      );
+    }
+  }
+
   return merge(
     {
-      oneOf: (
-        zodRef as z.ZodUnion<[z.ZodTypeAny, ...z.ZodTypeAny[]]>
-      )._def.options.map((schema) => generateSchema(schema, useOutput)),
+      oneOf: contents.map((schema) => generateSchema(schema, useOutput)),
     },
     zodRef.description ? { description: zodRef.description } : {},
     ...schemas