feat: add discriminator tag to generate if-then schemas (#1376)

src/AnnotationsReader/BasicAnnotationsReader.ts

@@ -19,6 +19,9 @@ export class BasicAnnotationsReader implements AnnotationsReader {
         "comment",
         "contentMediaType",
         "contentEncoding",
+
+        // Custom tag for if-then-else support.
+        "discriminator",
     ]);
     private static jsonTags = new Set<string>([
         "minimum",

src/Type/UnionType.ts

@@ -5,6 +5,7 @@ import { derefType } from "../Utils/derefType";
 
 export class UnionType extends BaseType {
     private readonly types: BaseType[];
+    private discriminator?: string = undefined;
 
     public constructor(types: readonly BaseType[]) {
         super();
@@ -20,6 +21,14 @@ export class UnionType extends BaseType {
         );
     }
 
+    public setDiscriminator(discriminator: string) {
+        this.discriminator = discriminator;
+    }
+
+    public getDiscriminator() {
+        return this.discriminator;
+    }
+
     public getId(): string {
         return `(${this.types.map((type) => type.getId()).join("|")})`;
     }

src/TypeFormatter/AnnotatedTypeFormatter.ts

@@ -2,7 +2,9 @@ import { Definition } from "../Schema/Definition";
 import { SubTypeFormatter } from "../SubTypeFormatter";
 import { AnnotatedType } from "../Type/AnnotatedType";
 import { BaseType } from "../Type/BaseType";
+import { UnionType } from "../Type/UnionType";
 import { TypeFormatter } from "../TypeFormatter";
+import { derefType } from "../Utils/derefType";
 
 export function makeNullable(def: Definition): Definition {
     const union: Definition[] | undefined = (def.oneOf as Definition[]) || def.anyOf;
@@ -50,6 +52,22 @@ export class AnnotatedTypeFormatter implements SubTypeFormatter {
         return type instanceof AnnotatedType;
     }
     public getDefinition(type: AnnotatedType): Definition {
+        const annotations = type.getAnnotations();
+
+        if ("discriminator" in annotations) {
+            const derefed = derefType(type.getType());
+            if (derefed instanceof UnionType) {
+                derefed.setDiscriminator(annotations.discriminator);
+                delete annotations.discriminator;
+            } else {
+                throw new Error(
+                    `Cannot assign discriminator tag to type: ${JSON.stringify(
+                        derefed
+                    )}. This tag can only be assigned to union types.`
+                );
+            }
+        }
+
         const def: Definition = {
             ...this.childTypeFormatter.getDefinition(type.getType()),
             ...type.getAnnotations(),

src/TypeFormatter/UnionTypeFormatter.ts

@@ -2,10 +2,12 @@ import { JSONSchema7 } from "json-schema";
 import { Definition } from "../Schema/Definition";
 import { SubTypeFormatter } from "../SubTypeFormatter";
 import { BaseType } from "../Type/BaseType";
+import { LiteralType } from "../Type/LiteralType";
 import { NeverType } from "../Type/NeverType";
 import { UnionType } from "../Type/UnionType";
 import { TypeFormatter } from "../TypeFormatter";
 import { derefType } from "../Utils/derefType";
+import { getTypeByKey } from "../Utils/typeKeys";
 import { uniqueArray } from "../Utils/uniqueArray";
 
 export class UnionTypeFormatter implements SubTypeFormatter {
@@ -20,6 +22,39 @@ export class UnionTypeFormatter implements SubTypeFormatter {
             .filter((item) => !(derefType(item) instanceof NeverType))
             .map((item) => this.childTypeFormatter.getDefinition(item));
 
+        const discriminator = type.getDiscriminator();
+        if (discriminator !== undefined) {
+            const kindTypes = type
+                .getTypes()
+                .filter((item) => !(derefType(item) instanceof NeverType))
+                .map((item) => getTypeByKey(item, new LiteralType(discriminator)));
+
+            const undefinedIndex = kindTypes.findIndex((item) => item === undefined);
+
+            if (undefinedIndex != -1) {
+                throw new Error(
+                    `Cannot find discriminator keyword "${discriminator}" in type ${JSON.stringify(
+                        type.getTypes()[undefinedIndex]
+                    )}.`
+                );
+            }
+
+            const kindDefinitions = kindTypes.map((item) => this.childTypeFormatter.getDefinition(item as BaseType));
+
+            const allOf = [];
+
+            for (let i = 0; i < definitions.length; i++) {
+                allOf.push({
+                    if: {
+                        properties: { [discriminator]: kindDefinitions[i] },
+                    },
+                    then: definitions[i],
+                });
+            }
+
+            return { allOf };
+        }
+
         // TODO: why is this not covered by LiteralUnionTypeFormatter?
         // special case for string literals | string -> string
         let stringType = true;

src/Utils/removeUnreachable.ts

@@ -58,6 +58,8 @@ function addReachable(
         } else if (items) {
             addReachable(items, definitions, reachable);
         }
+    } else if (definition.then) {
+        addReachable(definition.then, definitions, reachable);
     }
 }
 

test/invalid-data.test.ts

@@ -13,7 +13,7 @@ function assertSchema(name: string, type: string, message: string) {
             type: type,
             expose: "export",
             topRef: true,
-            jsDoc: "none",
+            jsDoc: "basic",
             skipTypeCheck: !!process.env.FAST_TEST,
         };
 
@@ -33,6 +33,28 @@ describe("invalid-data", () => {
 
     it("script-empty", assertSchema("script-empty", "MyType", `No root type "MyType" found`));
     it("duplicates", assertSchema("duplicates", "MyType", `Type "A" has multiple definitions.`));
+    it(
+        "missing-discriminator",
+        assertSchema(
+            "missing-discriminator",
+            "MyType",
+            'Cannot find discriminator keyword "type" in type ' +
+                '{"name":"B","type":{"id":"interface-1119825560-40-63-1119825560-0-124",' +
+                '"baseTypes":[],"properties":[],"additionalProperties":false,"nonPrimitive":false}}.'
+        )
+    );
+    it(
+        "non-union-discriminator",
+        assertSchema(
+            "non-union-discriminator",
+            "MyType",
+            "Cannot assign discriminator tag to type: " +
+                '{"id":"interface-2103469249-0-76-2103469249-0-77","baseTypes":[],' +
+                '"properties":[{"name":"name","type":{},"required":true}],' +
+                '"additionalProperties":false,"nonPrimitive":false}. ' +
+                "This tag can only be assigned to union types."
+        )
+    );
     it(
         "no-function-name",
         assertSchema(

test/invalid-data/missing-discriminator/main.ts

@@ -0,0 +1,10 @@
+export interface A {
+    type: string;
+}
+
+export interface B {}
+
+/**
+ * @discriminator type
+ */
+export type MyType = A | B;

test/invalid-data/non-union-discriminator/main.ts

@@ -0,0 +1,6 @@
+/**
+ * @discriminator name
+ */
+export interface MyType {
+    name: string;
+}

test/valid-data-annotations.test.ts

@@ -71,4 +71,6 @@ describe("valid-data-annotations", () => {
     it("annotation-ref", assertValidSchema("annotation-ref", "MyObject", "extended"));
 
     it("annotation-writeOnly", assertValidSchema("annotation-writeOnly", "MyObject", "basic"));
+
+    it("annotation-union-if-then", assertValidSchema("annotation-union-if-then", "Animal", "basic"));
 });

test/valid-data/annotation-union-if-then/main.ts

@@ -0,0 +1,14 @@
+export type Fish = {
+    animal_type: "fish";
+    found_in: "ocean" | "river";
+};
+
+export type Bird = {
+    animal_type: "bird";
+    can_fly: boolean;
+};
+
+/**
+ * @discriminator animal_type
+ */
+export type Animal = Bird | Fish;

test/valid-data/annotation-union-if-then/schema.json

@@ -0,0 +1,65 @@
+{
+  "$ref": "#/definitions/Animal",
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "definitions": {
+    "Animal": {
+      "allOf": [
+        {
+          "if": {
+            "properties": {
+              "animal_type": {
+                "const": "bird",
+                "type": "string"
+              }
+            }
+          },
+          "then": {
+            "$ref": "#/definitions/Bird"
+          }
+        },
+        {
+          "if": {
+            "properties": {
+              "animal_type": {
+                "const": "fish",
+                "type": "string"
+              }
+            }
+          },
+          "then": {
+            "$ref": "#/definitions/Fish"
+          }
+        }
+      ]
+    },
+    "Bird": {
+      "additionalProperties": false,
+      "properties": {
+        "animal_type": {
+          "const": "bird",
+          "type": "string"
+        },
+        "can_fly": {
+          "type": "boolean"
+        }
+      },
+      "required": ["animal_type", "can_fly"],
+      "type": "object"
+    },
+    "Fish": {
+      "additionalProperties": false,
+      "properties": {
+        "animal_type": {
+          "const": "fish",
+          "type": "string"
+        },
+        "found_in": {
+          "enum": ["ocean", "river"],
+          "type": "string"
+        }
+      },
+      "required": ["animal_type", "found_in"],
+      "type": "object"
+    }
+  }
+}