From 79999583a0dd4e4d9949c7959380b5e387786cee Mon Sep 17 00:00:00 2001
From: fmmoret <fmmoret2@gmail.com>
Date: Sat, 21 Oct 2023 13:41:41 -0700
Subject: [PATCH] force one of flag

---
 src/openapi-generator.ts | 12 ++++++++++++
 src/zod-extensions.ts    | 20 ++++++++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/openapi-generator.ts b/src/openapi-generator.ts
index 5af5dc8..dc737fb 100644
--- a/src/openapi-generator.ts
+++ b/src/openapi-generator.ts
@@ -960,6 +960,14 @@ export class OpenAPIGenerator {
         return this.generateSchemaWithRef(optionToGenerate);
       });
 
+      const internalOptions = this.getOptions(zodSchema);
+      if (internalOptions.forceOneOf) {
+        return {
+          oneOf: this.mapNullableOfArray(schemas, isNullable),
+          default: defaultValue,
+        };
+      }
+
       return {
         anyOf: this.mapNullableOfArray(schemas, isNullable),
         default: defaultValue,
@@ -1294,6 +1302,10 @@ export class OpenAPIGenerator {
     return this.getInternalMetadata(zodSchema)?.refId;
   }
 
+  private getOptions<T extends any>(zodSchema: ZodType<T>) {
+    return this.getInternalMetadata(zodSchema)?.options;
+  }
+
   private applySchemaMetadata(
     initialData: SchemaObject | ParameterObject | ReferenceObject,
     metadata: Partial<ZodOpenAPIMetadata>
diff --git a/src/zod-extensions.ts b/src/zod-extensions.ts
index 210728c..65f272f 100644
--- a/src/zod-extensions.ts
+++ b/src/zod-extensions.ts
@@ -27,6 +27,7 @@ export type ZodOpenAPIMetadata<T = any, E = ExampleValue<T>> = Omit<
 export interface ZodOpenAPIInternalMetadata {
   refId?: string;
   extendedFrom?: { refId: string; schema: ZodObject<ZodRawShape> };
+  options?: ZodToOpenAPIOptions;
 }
 
 export interface ZodOpenApiFullMetadata<T = any> {
@@ -34,6 +35,10 @@ export interface ZodOpenApiFullMetadata<T = any> {
   metadata?: ZodOpenAPIMetadata<T>;
 }
 
+export interface ZodToOpenAPIOptions {
+  forceOneOf?: boolean;
+}
+
 declare module 'zod' {
   interface ZodTypeDef {
     openapi?: ZodOpenApiFullMetadata;
@@ -54,6 +59,13 @@ declare module 'zod' {
       refId: string,
       metadata?: Partial<ZodOpenAPIMetadata<z.infer<T>>>
     ): T;
+
+    openapi<T extends ZodTypeAny>(
+      this: T,
+      refId: string,
+      metadata?: Partial<ZodOpenAPIMetadata<z.infer<T>>>,
+      options?: ZodToOpenAPIOptions,
+    ): T;
   }
 }
 
@@ -83,9 +95,12 @@ export function extendZodWithOpenApi(zod: typeof z) {
 
   zod.ZodType.prototype.openapi = function (
     refOrOpenapi: string | Partial<ZodOpenAPIMetadata<any>>,
-    metadata?: Partial<ZodOpenAPIMetadata<any>>
+    metadataOrOptions?: Partial<ZodOpenAPIMetadata<any>> | ZodToOpenAPIOptions,
+    options?: ZodToOpenAPIOptions,
   ) {
-    const openapi = typeof refOrOpenapi === 'string' ? metadata : refOrOpenapi;
+    const usingRef = typeof refOrOpenapi === 'string';
+    const openapi = (usingRef ? metadataOrOptions : refOrOpenapi) as Partial<ZodOpenAPIMetadata<any>> | undefined;
+    const resolvedOptions = (usingRef ? options : metadataOrOptions) as ZodToOpenAPIOptions | undefined;
 
     const { param, ...restOfOpenApi } = openapi ?? {};
 
@@ -94,6 +109,7 @@ export function extendZodWithOpenApi(zod: typeof z) {
       ...(typeof refOrOpenapi === 'string'
         ? { refId: refOrOpenapi }
         : undefined),
+      ...({ options: resolvedOptions } ?? {}),
     };
 
     const resultMetadata = {