fix: go back to generating enums out of unions containing enums (#2149)

Co-authored-by: Arthur Fiorette <[email protected]>

src/TypeFormatter/EnumTypeFormatter.ts

@@ -1,3 +1,4 @@
+import type { JSONSchema7TypeName } from "json-schema";
 import type { Definition } from "../Schema/Definition.js";
 import type { SubTypeFormatter } from "../SubTypeFormatter.js";
 import type { BaseType } from "../Type/BaseType.js";
@@ -17,11 +18,20 @@ export class EnumTypeFormatter implements SubTypeFormatter {
         // However, this formatter is used both for enum members and enum types,
         // so the side effect is that an enum type that contains just a single
         // value is represented as "const" too.
-        return values.length === 1
-            ? { type: types[0], const: values[0] }
-            : { type: types.length === 1 ? types[0] : types, enum: values };
+        return values.length === 1 ? { type: types[0], const: values[0] } : { type: toEnumType(types), enum: values };
     }
     public getChildren(type: EnumType): BaseType[] {
         return [];
     }
 }
+
+/**
+ * Unwraps the array if it contains only one type.
+ */
+export function toEnumType(types: JSONSchema7TypeName[]) {
+    if (types.length === 1) {
+        return types[0];
+    }
+
+    return types;
+}

src/TypeFormatter/LiteralUnionTypeFormatter.ts

@@ -2,84 +2,120 @@ import type { Definition } from "../Schema/Definition.js";
 import type { RawTypeName } from "../Schema/RawType.js";
 import type { SubTypeFormatter } from "../SubTypeFormatter.js";
 import type { BaseType } from "../Type/BaseType.js";
-import type { LiteralValue } from "../Type/LiteralType.js";
-import { LiteralType } from "../Type/LiteralType.js";
+import { EnumType } from "../Type/EnumType.js";
+import { LiteralType, type LiteralValue } from "../Type/LiteralType.js";
 import { NullType } from "../Type/NullType.js";
 import { StringType } from "../Type/StringType.js";
 import { UnionType } from "../Type/UnionType.js";
 import { typeName } from "../Utils/typeName.js";
-import { uniqueArray } from "../Utils/uniqueArray.js";
+import { toEnumType } from "./EnumTypeFormatter.js";
 
 export class LiteralUnionTypeFormatter implements SubTypeFormatter {
     public supportsType(type: BaseType): boolean {
         return type instanceof UnionType && type.getTypes().length > 0 && isLiteralUnion(type);
     }
-    public getDefinition(type: UnionType): Definition {
+
+    public getDefinition(unionType: UnionType): Definition {
         let hasString = false;
         let preserveLiterals = false;
         let allStrings = true;
         let hasNull = false;
 
-        const literals = type.getFlattenedTypes();
+        const literals = unionType.getFlattenedTypes();
 
         // filter out String types since we need to be more careful about them
-        const types = literals.filter((t) => {
-            if (t instanceof StringType) {
+        const types = literals.filter((literal) => {
+            if (literal instanceof StringType) {
                 hasString = true;
-                preserveLiterals = preserveLiterals || t.getPreserveLiterals();
+                preserveLiterals ||= literal.getPreserveLiterals();
                 return false;
-            } else if (t instanceof NullType) {
+            }
+
+            if (literal instanceof NullType) {
                 hasNull = true;
                 return true;
-            } else if (t instanceof LiteralType && !t.isString()) {
+            }
+
+            if (literal instanceof LiteralType && !literal.isString()) {
                 allStrings = false;
             }
 
             return true;
         });
 
         if (allStrings && hasString && !preserveLiterals) {
-            return {
-                type: hasNull ? ["string", "null"] : "string",
-            };
+            return hasNull ? { type: ["string", "null"] } : { type: "string" };
         }
 
-        const values = uniqueArray(types.map(getLiteralValue));
-        const typeNames = uniqueArray(types.map(getLiteralType));
+        const typeValues: Set<LiteralValue | null> = new Set();
+        const typeNames: Set<RawTypeName> = new Set();
 
-        const ret = {
-            type: typeNames.length === 1 ? typeNames[0] : typeNames,
-            enum: values,
-        };
-
-        if (preserveLiterals) {
-            return {
-                anyOf: [
-                    {
-                        type: "string",
-                    },
-                    ret,
-                ],
-            };
+        for (const type of types) {
+            appendTypeNames(type, typeNames);
+            appendTypeValues(type, typeValues);
         }
 
-        return ret;
+        const schema = {
+            type: toEnumType(Array.from(typeNames)),
+            enum: Array.from(typeValues),
+        };
+
+        return preserveLiterals ? { anyOf: [{ type: "string" }, schema] } : schema;
     }
-    public getChildren(type: UnionType): BaseType[] {
+
+    public getChildren(): BaseType[] {
         return [];
     }
 }
 
 export function isLiteralUnion(type: UnionType): boolean {
     return type
         .getFlattenedTypes()
-        .every((item) => item instanceof LiteralType || item instanceof NullType || item instanceof StringType);
+        .every(
+            (item) =>
+                item instanceof LiteralType ||
+                item instanceof NullType ||
+                item instanceof StringType ||
+                item instanceof EnumType,
+        );
 }
 
-function getLiteralValue(value: LiteralType | NullType): LiteralValue | null {
-    return value instanceof LiteralType ? value.getValue() : null;
+/**
+ * Appends all possible type names of a type to the given set.
+ */
+function appendTypeNames(type: BaseType, names: Set<RawTypeName>) {
+    if (type instanceof EnumType) {
+        for (const value of type.getValues()) {
+            names.add(typeName(value));
+        }
+
+        return;
+    }
+
+    if (type instanceof LiteralType) {
+        names.add(typeName(type.getValue()));
+        return;
+    }
+
+    names.add(typeName(null));
 }
 
-function getLiteralType(value: LiteralType | NullType): RawTypeName {
-    return value instanceof LiteralType ? typeName(value.getValue()) : "null";
+/**
+ * Appends all possible values of a type to the given set.
+ */
+function appendTypeValues(type: BaseType, values: Set<LiteralValue | null>) {
+    if (type instanceof EnumType) {
+        for (const value of type.getValues()) {
+            values.add(value);
+        }
+
+        return;
+    }
+
+    if (type instanceof LiteralType) {
+        values.add(type.getValue());
+        return;
+    }
+
+    values.add(null);
 }

src/Utils/uniqueArray.ts

@@ -1,9 +1,3 @@
 export function uniqueArray<T>(array: readonly T[]): T[] {
-    return array.reduce((result: T[], item: T) => {
-        if (!result.includes(item)) {
-            result.push(item);
-        }
-
-        return result;
-    }, []);
+    return array.filter((value, index) => array.indexOf(value) === index);
 }

test/valid-data-other.test.ts

@@ -9,6 +9,8 @@ describe("valid-data-other", () => {
     it("enums-mixed", assertValidSchema("enums-mixed", "Enum"));
     it("enums-member", assertValidSchema("enums-member", "MyObject"));
     it("enums-template-literal", assertValidSchema("enums-template-literal", "MyObject"));
+    it("enums-union", assertValidSchema("enums-union", "MyObject"));
+    it("exported-enums-union", assertValidSchema("exported-enums-union", "MyObject"));
 
     it("function-parameters-default-value", assertValidSchema("function-parameters-default-value", "myFunction"));
     it("function-parameters-declaration", assertValidSchema("function-parameters-declaration", "myFunction"));

test/valid-data/enums-union/main.ts

@@ -0,0 +1,26 @@
+enum Alphabet {
+    Alpha = "alpha",
+    Beta = "beta",
+    Omega = 666,
+}
+
+enum FileAccess {
+    None = 0,
+    Read = 1 << 1,
+    Write = 1 << 2,
+    ReadWrite = Read | Write,
+}
+
+export type MyObject = {
+    // All the members above should be output as enums, not anyOf
+    enumMembers: Alphabet.Alpha | Alphabet.Beta;
+    enumMemberWithLiteral: Alphabet.Alpha | "foo";
+    enumMemberWithLiteralAndNull: Alphabet.Alpha | "foo" | null;
+    enumMemberWithInterface: Alphabet.Alpha | { abc: string };
+    enumMembersWithNumber: Alphabet.Alpha | Alphabet.Omega;
+    wholeEnum: Alphabet; // Should output just all of Alphabet
+    wholeEnumWithLiteral: Alphabet | "bar"; // Should output all of Alphabet members (2 strings, 1 number) and "bar"
+    wholeEnumWithLiteralAndNull: Alphabet | "bar" | null;
+    twoDifferentEnumMembers: Alphabet.Alpha | FileAccess.Read;
+    twoDifferentWholeEnums: Alphabet | FileAccess;
+};

test/valid-data/enums-union/schema.json

@@ -0,0 +1,141 @@
+{
+  "$ref": "#/definitions/MyObject",
+  "$schema": "http://json-schema.org/draft-07/schema#",
+  "definitions": {
+    "MyObject": {
+      "additionalProperties": false,
+      "properties": {
+        "enumMemberWithInterface": {
+          "anyOf": [
+            {
+              "const": "alpha",
+              "type": "string"
+            },
+            {
+              "additionalProperties": false,
+              "properties": {
+                "abc": {
+                  "type": "string"
+                }
+              },
+              "required": [
+                "abc"
+              ],
+              "type": "object"
+            }
+          ]
+        },
+        "enumMemberWithLiteral": {
+          "enum": [
+            "alpha",
+            "foo"
+          ],
+          "type": "string"
+        },
+        "enumMemberWithLiteralAndNull": {
+          "enum": [
+            "alpha",
+            "foo",
+            null
+          ],
+          "type": [
+            "string",
+            "null"
+          ]
+        },
+        "enumMembers": {
+          "enum": [
+            "alpha",
+            "beta"
+          ],
+          "type": "string"
+        },
+        "enumMembersWithNumber": {
+          "enum": [
+            "alpha",
+            666
+          ],
+          "type": [
+            "string",
+            "number"
+          ]
+        },
+        "twoDifferentEnumMembers": {
+          "enum": [
+            "alpha",
+            2
+          ],
+          "type": [
+            "string",
+            "number"
+          ]
+        },
+        "twoDifferentWholeEnums": {
+          "enum": [
+            "alpha",
+            "beta",
+            666,
+            0,
+            2,
+            4,
+            6
+          ],
+          "type": [
+            "string",
+            "number"
+          ]
+        },
+        "wholeEnum": {
+          "enum": [
+            "alpha",
+            "beta",
+            666
+          ],
+          "type": [
+            "string",
+            "number"
+          ]
+        },
+        "wholeEnumWithLiteral": {
+          "enum": [
+            "alpha",
+            "beta",
+            666,
+            "bar"
+          ],
+          "type": [
+            "string",
+            "number"
+          ]
+        },
+        "wholeEnumWithLiteralAndNull": {
+          "enum": [
+            "alpha",
+            "beta",
+            666,
+            "bar",
+            null
+          ],
+          "type": [
+            "string",
+            "number",
+            "null"
+          ]
+        }
+      },
+      "required": [
+        "enumMembers",
+        "enumMemberWithLiteral",
+        "enumMemberWithLiteralAndNull",
+        "enumMemberWithInterface",
+        "enumMembersWithNumber",
+        "wholeEnum",
+        "wholeEnumWithLiteral",
+        "wholeEnumWithLiteralAndNull",
+        "twoDifferentEnumMembers",
+        "twoDifferentWholeEnums"
+      ],
+      "type": "object"
+    }
+  }
+}