From 46f4a257e5d3c4a85307b3b60607cda0c9b644e3 Mon Sep 17 00:00:00 2001
From: Sam Chung <samchungy@gmail.com>
Date: Thu, 9 Nov 2023 10:50:53 +1100
Subject: [PATCH] Fix Nullable Behaviour (#187)

---
 src/create/schema/metadata.ts              |  3 +-
 src/create/schema/parsers/nullable.test.ts | 30 ++++--------
 src/create/schema/parsers/nullable.ts      | 55 +++++++++++++---------
 src/openapi.ts                             |  7 +++
 4 files changed, 50 insertions(+), 45 deletions(-)

diff --git a/src/create/schema/metadata.ts b/src/create/schema/metadata.ts
index 59d6277..a32104a 100644
--- a/src/create/schema/metadata.ts
+++ b/src/create/schema/metadata.ts
@@ -1,10 +1,11 @@
+import { isReferenceObject } from '../../openapi';
 import type { oas31 } from '../../openapi3-ts/dist';
 
 export const enhanceWithMetadata = (
   schemaObject: oas31.SchemaObject | oas31.ReferenceObject,
   metadata: oas31.SchemaObject | oas31.ReferenceObject,
 ): oas31.SchemaObject | oas31.ReferenceObject => {
-  if ('$ref' in schemaObject) {
+  if (isReferenceObject(schemaObject)) {
     if (Object.values(metadata).every((val) => val === undefined)) {
       return schemaObject;
     }
diff --git a/src/create/schema/parsers/nullable.test.ts b/src/create/schema/parsers/nullable.test.ts
index 17a020c..cafbce0 100644
--- a/src/create/schema/parsers/nullable.test.ts
+++ b/src/create/schema/parsers/nullable.test.ts
@@ -25,16 +25,10 @@ describe('createNullableSchema', () => {
       expect(result).toStrictEqual(expected);
     });
 
-    it('creates an oneOf nullable schema for registered schemas', () => {
+    it('creates an allOf nullable schema for registered schemas', () => {
       const expected: oas30.SchemaObject = {
-        oneOf: [
-          {
-            $ref: '#/components/schemas/a',
-          },
-          {
-            nullable: true,
-          },
-        ],
+        allOf: [{ $ref: '#/components/schemas/a' }],
+        nullable: true,
       };
       const registered = z.string().openapi({ ref: 'a' });
       const schema = registered.optional().nullable();
@@ -65,10 +59,8 @@ describe('createNullableSchema', () => {
             },
             required: ['b'],
           },
-          {
-            nullable: true,
-          },
         ],
+        nullable: true,
       };
       const schema = z
         .union([z.object({ a: z.string() }), z.object({ b: z.string() })])
@@ -84,15 +76,11 @@ describe('createNullableSchema', () => {
         type: 'object',
         properties: {
           b: {
-            oneOf: [
-              {
-                allOf: [{ $ref: '#/components/schemas/a' }],
-                type: 'object',
-                properties: { b: { type: 'string' } },
-                required: ['b'],
-              },
-              { nullable: true },
-            ],
+            allOf: [{ $ref: '#/components/schemas/a' }],
+            type: 'object',
+            properties: { b: { type: 'string' } },
+            required: ['b'],
+            nullable: true,
           },
         },
         required: ['b'],
diff --git a/src/create/schema/parsers/nullable.ts b/src/create/schema/parsers/nullable.ts
index eeca210..6d50981 100644
--- a/src/create/schema/parsers/nullable.ts
+++ b/src/create/schema/parsers/nullable.ts
@@ -1,6 +1,6 @@
 import type { ZodNullable, ZodTypeAny } from 'zod';
 
-import { satisfiesVersion } from '../../../openapi';
+import { isReferenceObject, satisfiesVersion } from '../../../openapi';
 import type { oas31 } from '../../../openapi3-ts/dist';
 import type { ZodOpenApiVersion } from '../../document';
 import { type SchemaState, createSchemaObject } from '../../schema';
@@ -8,44 +8,53 @@ import { type SchemaState, createSchemaObject } from '../../schema';
 export const createNullableSchema = <T extends ZodTypeAny>(
   zodNullable: ZodNullable<T>,
   state: SchemaState,
-): oas31.SchemaObject => {
+): oas31.SchemaObject | oas31.ReferenceObject => {
   const schemaObject = createSchemaObject(zodNullable.unwrap(), state, [
     'nullable',
   ]);
 
-  if ('$ref' in schemaObject || schemaObject.allOf) {
-    return {
-      oneOf: mapNullOf([schemaObject], state.components.openapi),
-    };
-  }
+  if (satisfiesVersion(state.components.openapi, '3.1.0')) {
+    if (isReferenceObject(schemaObject) || schemaObject.allOf) {
+      return {
+        oneOf: mapNullOf([schemaObject], state.components.openapi),
+      };
+    }
+
+    if (schemaObject.oneOf) {
+      const { oneOf, ...schema } = schemaObject;
+      return {
+        oneOf: mapNullOf(oneOf, state.components.openapi),
+        ...schema,
+      };
+    }
+
+    if (schemaObject.anyOf) {
+      const { anyOf, ...schema } = schemaObject;
+      return {
+        anyOf: mapNullOf(anyOf, state.components.openapi),
+        ...schema,
+      };
+    }
+
+    const { type, ...schema } = schemaObject;
 
-  if (schemaObject.oneOf) {
-    const { oneOf, ...schema } = schemaObject;
     return {
-      oneOf: mapNullOf(oneOf, state.components.openapi),
+      type: mapNullType(type),
       ...schema,
     };
   }
 
-  if (schemaObject.anyOf) {
-    const { anyOf, ...schema } = schemaObject;
+  if (isReferenceObject(schemaObject)) {
     return {
-      anyOf: mapNullOf(anyOf, state.components.openapi),
-      ...schema,
-    };
+      allOf: [schemaObject],
+      nullable: true,
+    } as oas31.SchemaObject;
   }
 
   const { type, ...schema } = schemaObject;
 
-  if (satisfiesVersion(state.components.openapi, '3.1.0')) {
-    return {
-      type: mapNullType(type),
-      ...schema,
-    };
-  }
-
   return {
-    type,
+    ...(type && { type }),
     nullable: true,
     ...schema,
     // https://github.com/OAI/OpenAPI-Specification/blob/main/proposals/2019-10-31-Clarify-Nullable.md#if-a-schema-specifies-nullable-true-and-enum-1-2-3-does-that-schema-allow-null-values-see-1900
diff --git a/src/openapi.ts b/src/openapi.ts
index 79fa89a..03e425f 100644
--- a/src/openapi.ts
+++ b/src/openapi.ts
@@ -1,3 +1,5 @@
+import type { oas31 } from './openapi3-ts/dist';
+
 export const openApiVersions = [
   '3.0.0',
   '3.0.1',
@@ -12,3 +14,8 @@ export const satisfiesVersion = (
   test: OpenApiVersion,
   against: OpenApiVersion,
 ) => openApiVersions.indexOf(test) >= openApiVersions.indexOf(against);
+
+export const isReferenceObject = (
+  schemaOrRef: oas31.SchemaObject | oas31.ReferenceObject,
+): schemaOrRef is oas31.ReferenceObject =>
+  Boolean('$ref' in schemaOrRef && schemaOrRef.$ref);