Add support for ZodLazy, ZodPipeline, ZodBrand and ZodEffects for par…

…ameter schemas (#229) Co-authored-by: Andy Edwards <[email protected]>
samchungy · Mar 1, 2024 · 468b1fa · 468b1fa
1 parent f8ee426
commit 468b1fa
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 9 deletions.
diff --git a/src/create/document.ts b/src/create/document.ts
@@ -1,4 +1,4 @@
-import type { AnyZodObject, ZodType } from 'zod';
+import type { AnyZodObject, ZodType, ZodTypeDef } from 'zod';
 
 import type { OpenApiVersion } from '../openapi';
 import type { oas30, oas31 } from '../openapi3-ts/dist';
@@ -46,7 +46,8 @@ export interface ZodOpenApiResponsesObject
 }
 
 export type ZodOpenApiParameters = {
-  [type in oas31.ParameterLocation & oas30.ParameterLocation]?: AnyZodObject;
+  [type in oas31.ParameterLocation &
+    oas30.ParameterLocation]?: ZodObjectInputType;
 };
 
 export interface ZodOpenApiOperationObject
@@ -132,6 +133,12 @@ export interface ZodOpenApiObject
   components?: ZodOpenApiComponentsObject;
 }
 
+export type ZodObjectInputType<
+  Output = unknown,
+  Def extends ZodTypeDef = ZodTypeDef,
+  Input = Record<string, unknown>,
+> = ZodType<Output, Def, Input>;
+
 export const createDocument = (
   zodOpenApiObject: ZodOpenApiObject,
 ): oas31.OpenAPIObject => {

diff --git a/src/create/parameters.test.ts b/src/create/parameters.test.ts
@@ -267,4 +267,74 @@ describe('createParametersObject', () => {
 
     expect(result).toStrictEqual(expectedResult);
   });
+
+  it('should support wrapped ZodObject schema', () => {
+    const expectedResult: oas31.OperationObject['parameters'] = [
+      {
+        in: 'path',
+        name: 'b',
+        schema: {
+          type: 'string',
+        },
+        required: true,
+      },
+      {
+        in: 'query',
+        name: 'a',
+        schema: {
+          type: 'string',
+        },
+        required: true,
+      },
+      {
+        in: 'cookie',
+        name: 'c',
+        schema: {
+          type: 'string',
+        },
+        required: true,
+      },
+      {
+        in: 'header',
+        name: 'd',
+        schema: {
+          type: 'string',
+        },
+        required: true,
+      },
+    ];
+    const result = createParametersObject(
+      [],
+      {
+        path: z
+          .lazy(() =>
+            z.object({ b: z.string() }).pipe(z.object({ b: z.string() })),
+          )
+          .brand('test')
+          .transform((x) => x),
+        query: z
+          .lazy(() =>
+            z.object({ a: z.string() }).pipe(z.object({ a: z.string() })),
+          )
+          .brand('test')
+          .transform((x) => x),
+        cookie: z
+          .lazy(() =>
+            z.object({ c: z.string() }).pipe(z.object({ c: z.string() })),
+          )
+          .brand('test')
+          .transform((x) => x),
+        header: z
+          .lazy(() =>
+            z.object({ d: z.string() }).pipe(z.object({ d: z.string() })),
+          )
+          .brand('test')
+          .transform((x) => x),
+      },
+      getDefaultComponents(),
+      ['parameters'],
+    );
+
+    expect(result).toStrictEqual(expectedResult);
+  });
 });
diff --git a/src/create/parameters.ts b/src/create/parameters.ts
@@ -1,10 +1,10 @@
 import type { AnyZodObject, ZodRawShape, ZodType } from 'zod';
 
 import type { oas30, oas31 } from '../openapi3-ts/dist';
-import { isAnyZodType } from '../zodType';
+import { isAnyZodType, isZodType } from '../zodType';
 
 import type { ComponentsObject } from './components';
-import type { ZodOpenApiParameters } from './document';
+import type { ZodObjectInputType, ZodOpenApiParameters } from './document';
 import { type SchemaState, createSchema } from './schema';
 import { isOptionalSchema } from './schema/parsers/optional';
 
@@ -101,17 +101,18 @@ export const createParamOrRef = (
 
 const createParameters = (
   type: keyof ZodOpenApiParameters,
-  zodObject: AnyZodObject | undefined,
+  zodObjectType: ZodObjectInputType | undefined,
   components: ComponentsObject,
   subpath: string[],
 ): Array<oas31.ParameterObject | oas31.ReferenceObject> => {
-  if (!zodObject) {
+  if (!zodObjectType) {
     return [];
   }
 
-  return Object.entries(zodObject.shape as ZodRawShape).map(
-    ([key, zodSchema]: [string, ZodType]) =>
-      createParamOrRef(zodSchema, components, [...subpath, key], type, key),
+  const zodObject = getZodObject(zodObjectType, 'input').shape as ZodRawShape;
+
+  return Object.entries(zodObject).map(([key, zodSchema]: [string, ZodType]) =>
+    createParamOrRef(zodSchema, components, [...subpath, key], type, key),
   );
 };
 
@@ -200,3 +201,29 @@ export const createParametersObject = (
 
   return combinedParameters.length ? combinedParameters : undefined;
 };
+
+export const getZodObject = (
+  schema: ZodObjectInputType,
+  type: 'input' | 'output',
+): AnyZodObject => {
+  if (isZodType(schema, 'ZodObject')) {
+    return schema;
+  }
+  if (isZodType(schema, 'ZodLazy')) {
+    return getZodObject(schema.schema as ZodObjectInputType, type);
+  }
+  if (isZodType(schema, 'ZodEffects')) {
+    return getZodObject(schema.innerType() as ZodObjectInputType, type);
+  }
+  if (isZodType(schema, 'ZodBranded')) {
+    return getZodObject(schema.unwrap() as ZodObjectInputType, type);
+  }
+  if (isZodType(schema, 'ZodPipeline')) {
+    if (type === 'input') {
+      return getZodObject(schema._def.in as ZodObjectInputType, type);
+    }
+    return getZodObject(schema._def.out as ZodObjectInputType, type);
+  }
+
+  throw new Error('failed to find ZodObject in schema');
+};