Discriminator Mapping Support for oneOf composition (#1574)

* Added discriminator mapping support for oneOf composition

* Discriminator support proposal cleanup

* Refactor discriminator implementation to handle all oneOf+mapping combinations

* Refactor global context to a single property for all discriminator values

* Update examples with new discriminator output

* Support discriminator enum mappings using the same ref multiple times

* Updated examples with improved discriminator support

---------

Co-authored-by: mat <[email protected]>

Author: mzronek

packages/openapi-typescript/examples/digital-ocean-api.ts

@@ -9565,11 +9565,10 @@ export interface components {
         /** @description Specifies the action that will be taken on the Droplet. */
         droplet_action: {
             /**
-             * @description The type of action to initiate for the Droplet.
-             * @example reboot
+             * @description The type of action to initiate for the Droplet. (enum property replaced by openapi-typescript)
              * @enum {string}
              */
-            type: "enable_backups" | "disable_backups" | "reboot" | "power_cycle" | "shutdown" | "power_off" | "power_on" | "restore" | "password_reset" | "resize" | "rebuild" | "rename" | "change_kernel" | "enable_ipv6" | "snapshot";
+            type: "enable_backups" | "disable_backups" | "power_cycle" | "shutdown" | "power_off" | "power_on" | "enable_ipv6";
         };
         droplet_action_restore: components["schemas"]["droplet_action"] & {
             /**
@@ -9617,6 +9616,12 @@ export interface components {
              * @example Nifty New Snapshot
              */
             name?: string;
+        } & {
+            /**
+             * @description discriminator enum property added by openapi-typescript
+             * @enum {string}
+             */
+            type: "snapshot";
         };
         firewall_rule_base: {
             /**

packages/openapi-typescript/src/index.ts

@@ -71,7 +71,7 @@ export default async function openapiTS(
     additionalProperties: options.additionalProperties ?? false,
     alphabetize: options.alphabetize ?? false,
     defaultNonNullable: options.defaultNonNullable ?? true,
-    discriminators: scanDiscriminators(schema),
+    discriminators: scanDiscriminators(schema, options),
     emptyObjectsUnknown: options.emptyObjectsUnknown ?? false,
     enum: options.enum ?? false,
     excludeDeprecated: options.excludeDeprecated ?? false,

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 } 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) {
@@ -165,15 +171,154 @@ export function resolveRef<T>(
   return node;
 }
 
+function createDiscriminatorEnum(
+  values: string[],
+  prevSchema?: SchemaObject,
+): SchemaObject {
+  return {
+    type: "string",
+    enum: values,
+    description: prevSchema?.description
+      ? `${prevSchema.description} (enum property replaced by openapi-typescript)`
+      : `discriminator enum property added by openapi-typescript`,
+  };
+}
+
+type InternalDiscriminatorMapping = Record<
+  string,
+  { inferred?: string; defined?: string[] }
+>;
+
 /** Return a key–value map of discriminator objects found in a schema */
-export function scanDiscriminators(schema: OpenAPI3) {
-  const discriminators: Record<string, DiscriminatorObject> = {};
+export function scanDiscriminators(
+  schema: OpenAPI3,
+  options: OpenAPITSOptions,
+) {
+  // all discriminator objects found in the schema
+  const objects: Record<string, DiscriminatorObject> = {};
+
+  // refs of all mapped schema objects we have successfully handled to infer the discriminator enum value
+  const refsHandled: 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) => {
-    if ((obj?.discriminator as DiscriminatorObject)?.propertyName) {
-      discriminators[createRef(path)] =
-        obj.discriminator as DiscriminatorObject;
+    const discriminator = obj?.discriminator as DiscriminatorObject | undefined;
+    if (!discriminator?.propertyName) {
+      return;
+    }
+
+    // collect discriminator object for later usage
+    const ref = createRef(path);
+
+    objects[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: InternalDiscriminatorMapping = {};
+
+    // 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 inferred discriminator enum value
+        const value = item.$ref.split("/").pop();
+
+        if (value) {
+          if (!mapping[item.$ref]) {
+            mapping[item.$ref] = { inferred: value };
+          } else {
+            mapping[item.$ref].inferred = 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];
+        if (!mappedRef) {
+          continue;
+        }
+
+        if (!mapping[mappedRef]?.defined) {
+          // this overrides inferred values, but we don't need them anymore as soon as we have a defined value
+          mapping[mappedRef] = { defined: [] };
+        }
+
+        mapping[mappedRef].defined?.push(mappedValue);
+      }
+    }
+
+    for (const [mappedRef, { inferred, defined }] of Object.entries(mapping)) {
+      if (refsHandled.includes(mappedRef)) {
+        continue;
+      }
+
+      if (!inferred && !defined) {
+        continue;
+      }
+
+      // prefer defined values over automatically inferred ones
+      // the inferred enum values from the schema might not represent the actual enum values of the discriminator,
+      // so if we have defined values, use them instead
+      const mappedValues = defined ?? [inferred!];
+      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",
+          // discriminator enum properties always need to be required
+          required: [discriminator.propertyName],
+          properties: {
+            [discriminator.propertyName]: createDiscriminatorEnum(mappedValues),
+          },
+        });
+
+        refsHandled.push(mappedRef);
+      } else if (
+        typeof resolvedSchema === "object" &&
+        "type" in resolvedSchema &&
+        resolvedSchema.type === "object"
+      ) {
+        // if the schema is an object, we can apply the discriminator enums to its properties
+        if (!resolvedSchema.properties) {
+          resolvedSchema.properties = {};
+        }
+
+        // discriminator enum properties always need to be required
+        if (!resolvedSchema.required) {
+          resolvedSchema.required = [discriminator.propertyName];
+        } else if (
+          !resolvedSchema.required.includes(discriminator.propertyName)
+        ) {
+          resolvedSchema.required.push(discriminator.propertyName);
+        }
+
+        // add/replace the discriminator enum property
+        resolvedSchema.properties[discriminator.propertyName] =
+          createDiscriminatorEnum(
+            mappedValues,
+            resolvedSchema.properties[discriminator.propertyName],
+          );
+
+        refsHandled.push(mappedRef);
+      } else {
+        warn(
+          `Discriminator mapping has an invalid schema (neither an object schema nor an allOf array): ${mappedRef} => ${mappedValues.join(
+            ", ",
+          )} (Discriminator: ${ref})`,
+          options.silent,
+        );
+        continue;
+      }
     }
   });
 
@@ -185,19 +330,20 @@ export function scanDiscriminators(schema: OpenAPI3) {
       if (obj && Array.isArray(obj[key])) {
         for (const item of (obj as any)[key]) {
           if ("$ref" in item) {
-            if (discriminators[item.$ref]) {
-              discriminators[createRef(path)] = {
-                ...discriminators[item.$ref],
+            if (objects[item.$ref]) {
+              objects[createRef(path)] = {
+                ...objects[item.$ref],
               };
             }
           } else if (item.discriminator?.propertyName) {
-            discriminators[createRef(path)] = { ...item.discriminator };
+            objects[createRef(path)] = { ...item.discriminator };
           }
         }
       }
     }
   });
-  return discriminators;
+
+  return { objects, refsHandled };
 }
 
 /** Walk through any JSON-serializable (i.e. non-circular) object */

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.discriminators.refsHandled.includes(item.$ref)) {
+            const validRequired = (required ?? []).filter(
+              (key) => !!resolved.properties![key],
             );
+            if (validRequired.length) {
+              itemType = tsWithRequired(
+                itemType,
+                validRequired,
+                options.ctx.injectFooter,
+              );
+            }
           }
         }
       }
@@ -182,7 +185,7 @@ export function transformSchemaObjectWithComposition(
         );
       }
       const discriminator =
-        ("$ref" in item && options.ctx.discriminators[item.$ref]) ||
+        ("$ref" in item && options.ctx.discriminators.objects[item.$ref]) ||
         (item as any).discriminator; // eslint-disable-line @typescript-eslint/no-explicit-any
       if (discriminator) {
         output.push(tsOmit(itemType, [discriminator.propertyName]));
@@ -414,7 +417,8 @@ function transformSchemaObjectCore(
     // ctx.discriminators. But stop objects from referencing their own
     // discriminator meant for children (!schemaObject.discriminator)
     const discriminator =
-      !schemaObject.discriminator && options.ctx.discriminators[options.path!];
+      !schemaObject.discriminator &&
+      options.ctx.discriminators.objects[options.path!];
     if (discriminator) {
       coreObjectType.unshift(
         createDiscriminatorProperty(discriminator, {

packages/openapi-typescript/src/types.ts

@@ -686,7 +686,10 @@ export interface GlobalContext {
   additionalProperties: boolean;
   alphabetize: boolean;
   defaultNonNullable: boolean;
-  discriminators: Record<string, DiscriminatorObject>;
+  discriminators: {
+    objects: Record<string, DiscriminatorObject>;
+    refsHandled: string[];
+  };
   emptyObjectsUnknown: boolean;
   enum: boolean;
   excludeDeprecated: boolean;