Refactor discriminator implementation to handle all oneOf+mapping combinations

mzronek · mzronek · commit b08447243ba6 · 2024-03-03T13:19:49.000+01:00
diff --git a/packages/openapi-typescript/src/index.ts b/packages/openapi-typescript/src/index.ts
@@ -67,11 +67,17 @@ export default async function openapiTS(
     silent: options.silent ?? false,
   });
 
+  const { discriminators, discriminatorRefsHandled } = scanDiscriminators(
+    schema,
+    options,
+  );
+
   const ctx: GlobalContext = {
     additionalProperties: options.additionalProperties ?? false,
     alphabetize: options.alphabetize ?? false,
     defaultNonNullable: options.defaultNonNullable ?? true,
-    discriminators: scanDiscriminators(schema, options),
+    discriminators,
+    discriminatorRefsHandled,
     emptyObjectsUnknown: options.emptyObjectsUnknown ?? false,
     enum: options.enum ?? false,
     excludeDeprecated: options.excludeDeprecated ?? false,
diff --git a/packages/openapi-typescript/src/lib/utils.ts b/packages/openapi-typescript/src/lib/utils.ts
@@ -5,7 +5,13 @@ import {
 import c from "ansi-colors";
 import supportsColor from "supports-color";
 import ts from "typescript";
-import type { DiscriminatorObject, OpenAPI3, OpenAPITSOptions, SchemaObject } from "../types.js";
+import type {
+  DiscriminatorObject,
+  OpenAPI3,
+  OpenAPITSOptions,
+  ReferenceObject,
+  SchemaObject,
+} from "../types.js";
 import { tsLiteral, tsModifiers, tsPropertyIndex } from "./ts.js";
 
 if (!supportsColor.stdout || supportsColor.stdout.hasBasic === false) {
@@ -174,42 +180,103 @@ function createDiscriminatorEnum(value: string): SchemaObject {
 }
 
 /** Return a key–value map of discriminator objects found in a schema */
-export function scanDiscriminators(schema: OpenAPI3, options: OpenAPITSOptions) {
+export function scanDiscriminators(
+  schema: OpenAPI3,
+  options: OpenAPITSOptions,
+) {
   const discriminators: Record<string, DiscriminatorObject> = {};
-  const enumsAppended: string[] = [];
+  // discriminator objects which we have successfully handled to infer the discriminator enum value
+  const discriminatorRefsHandled: string[] = [];
 
-  // perform 2 passes: first, collect all discriminator definitions
+  // perform 2 passes: first, collect all discriminator definitions and handle oneOf and mappings
   walk(schema, (obj, path) => {
     const discriminator = obj?.discriminator as DiscriminatorObject | undefined;
-    if (discriminator?.propertyName) {
-      const ref = createRef(path);
-
-      discriminators[ref] = discriminator;
-
-      // if a mapping is available we will help Typescript by adding the discriminator enum with its mapped value to each schema
-      if (discriminator.mapping) {
-        for (const [mappedValue, mappedRef] of Object.entries(
-          discriminator.mapping
-        )) {
-          if (enumsAppended.includes(mappedRef)) {
-            continue;
-          }
+    if (!discriminator?.propertyName) {
+      return;
+    }
 
-          const resolved = resolveRef<SchemaObject>(schema, mappedRef, { silent: options.silent ?? false });
-          if (resolved?.allOf) {
-            resolved.allOf.push({ type: "object", properties: { [discriminator.propertyName]: createDiscriminatorEnum(mappedValue) } });
-            enumsAppended.push(mappedRef);
-          } else if (typeof resolved === 'object' && 'type' in resolved && resolved.type === 'object') {
-            if (!resolved.properties) {
-              resolved.properties = {};
-            }
+    // add to discriminators object for later usage
+    const ref = createRef(path);
 
-            resolved.properties[discriminator.propertyName] = createDiscriminatorEnum(mappedValue);
-            enumsAppended.push(mappedRef);
-          } else {
-            continue;
-          }
+    discriminators[ref] = discriminator;
+
+    // if a mapping is available we will help Typescript to infer properties by adding the discriminator enum with its single mapped value to each schema
+    // we only handle the mapping in advance for discriminator + oneOf compositions right now
+    if (!obj?.oneOf || !Array.isArray(obj.oneOf)) {
+      return;
+    }
+
+    const oneOf: (SchemaObject | ReferenceObject)[] = obj.oneOf;
+    const mapping: Record<string, string> = {};
+
+    // the mapping can be inferred from the oneOf refs next to the discriminator object
+    for (const item of oneOf) {
+      if ("$ref" in item) {
+        // the name of the schema is the infered discriminator enum value
+        const value = item.$ref.split("/").pop();
+
+        if (value) {
+          mapping[item.$ref] = value;
+        }
+      }
+    }
+
+    // the mapping can be defined in the discriminator object itself
+    if (discriminator.mapping) {
+      for (const mappedValue in discriminator.mapping) {
+        const mappedRef = discriminator.mapping[mappedValue];
+
+        mapping[mappedRef] = mappedValue;
+      }
+    }
+
+    for (const [mappedRef, mappedValue] of Object.entries(mapping)) {
+      if (discriminatorRefsHandled.includes(mappedRef)) {
+        continue;
+      }
+
+      const resolvedSchema = resolveRef<SchemaObject>(schema, mappedRef, {
+        silent: options.silent ?? false,
+      });
+      if (resolvedSchema?.allOf) {
+        // if the schema is an allOf, we can append a new schema object to the allOf array
+        resolvedSchema.allOf.push({
+          type: "object",
+          required: [discriminator.propertyName],
+          properties: {
+            [discriminator.propertyName]: createDiscriminatorEnum(mappedValue),
+          },
+        });
+
+        discriminatorRefsHandled.push(mappedRef);
+      } else if (
+        typeof resolvedSchema === "object" &&
+        "type" in resolvedSchema &&
+        resolvedSchema.type === "object"
+      ) {
+        // if the schema is an object, we can add/replace the discriminator enum to it
+        if (!resolvedSchema.properties) {
+          resolvedSchema.properties = {};
+        }
+
+        if (!resolvedSchema.required) {
+          resolvedSchema.required = [discriminator.propertyName];
+        } else if (
+          !resolvedSchema.required.includes(discriminator.propertyName)
+        ) {
+          resolvedSchema.required.push(discriminator.propertyName);
         }
+
+        resolvedSchema.properties[discriminator.propertyName] =
+          createDiscriminatorEnum(mappedValue);
+
+        discriminatorRefsHandled.push(mappedRef);
+      } else {
+        warn(
+          `Discriminator mapping has an invalid schema (neither an object schema nor an allOf array): ${mappedRef} => ${mappedValue} (Discriminator: ${ref})`,
+          options.silent,
+        );
+        continue;
       }
     }
   });
@@ -234,7 +301,8 @@ export function scanDiscriminators(schema: OpenAPI3, options: OpenAPITSOptions)
       }
     }
   });
-  return discriminators;
+
+  return { discriminators, discriminatorRefsHandled };
 }
 
 /** Walk through any JSON-serializable (i.e. non-circular) object */
diff --git a/packages/openapi-typescript/src/transform/schema-object.ts b/packages/openapi-typescript/src/transform/schema-object.ts
@@ -157,16 +157,19 @@ export function transformSchemaObjectWithComposition(
           typeof resolved === "object" &&
           "properties" in resolved
         ) {
-          // don’t try and make keys required if the $ref doesn’t have them
-          const validRequired = (required ?? []).filter(
-            (key) => !!resolved.properties![key],
-          );
-          if (validRequired.length) {
-            itemType = tsWithRequired(
-              itemType,
-              validRequired,
-              options.ctx.injectFooter,
+          // don’t try and make keys required if we have already handled the item (discriminator property was already added as required)
+          // or the $ref doesn’t have them
+          if (!options.ctx.discriminatorRefsHandled.includes(item.$ref)) {
+            const validRequired = (required ?? []).filter(
+              (key) => !!resolved.properties![key],
             );
+            if (validRequired.length) {
+              itemType = tsWithRequired(
+                itemType,
+                validRequired,
+                options.ctx.injectFooter,
+              );
+            }
           }
         }
       }
diff --git a/packages/openapi-typescript/src/types.ts b/packages/openapi-typescript/src/types.ts
@@ -687,6 +687,7 @@ export interface GlobalContext {
   alphabetize: boolean;
   defaultNonNullable: boolean;
   discriminators: Record<string, DiscriminatorObject>;
+  discriminatorRefsHandled: string[];
   emptyObjectsUnknown: boolean;
   enum: boolean;
   excludeDeprecated: boolean;
diff --git a/packages/openapi-typescript/test/discriminators.test.ts b/packages/openapi-typescript/test/discriminators.test.ts
@@ -195,7 +195,7 @@ export type operations = Record<string, never>;`,
       },
     ],
     [
-      "oneOf",
+      "oneOf > implicit mapping",
       {
         given: {
           openapi: "3.1",
@@ -246,12 +246,27 @@ export interface components {
         Pet: components["schemas"]["Cat"] | components["schemas"]["Dog"] | components["schemas"]["Lizard"];
         Cat: {
             name?: string;
+            /**
+             * @description discriminator enum property added by openapi-typescript
+             * @enum {string}
+             */
+            petType: "Cat";
         };
         Dog: {
             bark?: string;
+            /**
+             * @description discriminator enum property added by openapi-typescript
+             * @enum {string}
+             */
+            petType: "dog";
         };
         Lizard: {
             lovesRocks?: boolean;
+            /**
+             * @description discriminator enum property added by openapi-typescript
+             * @enum {string}
+             */
+            petType: "Lizard";
         };
     };
     responses: never;
@@ -266,7 +281,7 @@ export type operations = Record<string, never>;`,
       },
     ],
     [
-      "oneOf > mapping support > replace discriminator enum",
+      "oneOf > explicit mapping > replace discriminator enum",
       {
         given: {
           openapi: "3.1.0",
@@ -409,7 +424,7 @@ export type operations = Record<string, never>;`,
       },
     ],
     [
-      "oneOf > mapping support > append enum in allOf",
+      "oneOf > explicit mapping > append enum in allOf",
       {
         given: {
           openapi: "3.1.0",
@@ -547,7 +562,7 @@ export interface components {
              * @description discriminator enum property added by openapi-typescript
              * @enum {string}
              */
-            type?: "simple";
+            type: "simple";
         };
         complexObject: components["schemas"]["baseObject"] & {
             complex?: boolean;
@@ -556,7 +571,7 @@ export interface components {
              * @description discriminator enum property added by openapi-typescript
              * @enum {string}
              */
-            type?: "complex";
+            type: "complex";
         };
         /** @enum {string} */
         type: "simple" | "complex";
diff --git a/packages/openapi-typescript/test/test-helpers.ts b/packages/openapi-typescript/test/test-helpers.ts
@@ -11,6 +11,7 @@ export const DEFAULT_CTX: GlobalContext = {
   arrayLength: false,
   defaultNonNullable: true,
   discriminators: {},
+  discriminatorRefsHandled: [],
   emptyObjectsUnknown: false,
   enum: false,
   excludeDeprecated: false,
