Merge pull request #32 from code0-tech/feat/recursive-type-schemas

Fix recursive type schemas

Author: nicosammito

ts/src/internal/zod-schema.ts

@@ -30,15 +30,34 @@ function buildOverrides(skip: ZodTypeAny): TypeOverrideMap {
 }
 
 export function zodToTypeString(schema: ZodTypeAny): string {
+    const selfIdentifier = schemaRegistry.get(schema);
+    const auxiliaryTypeStore = createAuxiliaryTypeStore();
+    let rootVisited = false;
     const {node} = zodToTs(schema, {
-        auxiliaryTypeStore: createAuxiliaryTypeStore(),
+        auxiliaryTypeStore,
         overrides: buildOverrides(schema),
+        overrideFunction: (candidate, typescript) => {
+            if (candidate !== schema || selfIdentifier === undefined) return undefined;
+            if (!rootVisited) {
+                rootVisited = true;
+                return undefined;
+            }
+            return typescript.factory.createTypeReferenceNode(typescript.factory.createIdentifier(selfIdentifier));
+        },
     });
-    return printNode(node, {removeComments: true, omitTrailingSemicolon: true}).replace(/\s+/g, " ").trim();
+    const type = printNode(node, {removeComments: true, omitTrailingSemicolon: true}).replace(/\s+/g, " ").trim();
+    if (auxiliaryTypeStore.definitions.size > 0) {
+        const subject = selfIdentifier ? `data type "${selfIdentifier}"` : "the schema";
+        throw new Error(
+            `Cannot generate a type string for ${subject}: it contains recursive schemas that cannot be inlined. ` +
+            `Register each recursive schema as its own data type so it can be referenced by identifier.`
+        );
+    }
+    return type;
 }
 
 export function zodToRules(schema: ZodTypeAny): DefinitionDataTypeRule[] {
-    const json = toJSONSchema(schema) as JsonSchema;
+    const json = toJSONSchema(schema, {cycles: "ref"}) as JsonSchema;
     const rules: DefinitionDataTypeRule[] = [];
     if (json.pattern) {
         rules.push({config: {oneofKind: "regex", regex: {pattern: json.pattern}}});

ts/src/map/datatype.map.ts

@@ -18,7 +18,9 @@ export const dataTypeMap = (klass: DataTypeClass): DataTypeProps => {
 
     return {
         identifier,
-        type: zodToTypeString(schema),
+        get type() {
+            return zodToTypeString(schema)
+        },
         rules: zodToRules(schema),
         name,
         displayMessage,

ts/test/datatype.test.ts

@@ -0,0 +1,83 @@
+import "reflect-metadata";
+import { describe, expect, it } from "vitest";
+import { z, type ZodTypeAny } from "zod";
+import { Identifier } from "../src/decorators/meta.dec";
+import { Schema } from "../src/decorators/datatype.dec";
+import { dataTypeMap } from "../src/map/datatype.map";
+
+describe("dataTypeMap type generation", () => {
+    it("inlines plain schemas", () => {
+        @Identifier("PlainType")
+        @Schema(z.object({ name: z.string() }))
+        class PlainType {}
+
+        expect(dataTypeMap(PlainType).type).toBe("{ name: string; }");
+    });
+
+    it("resolves mutually recursive schemas by identifier regardless of registration order", () => {
+        const OrderSchema: ZodTypeAny = z.lazy(() =>
+            z.object({ addresses: z.array(AddressSchema) })
+        );
+        const AddressSchema: ZodTypeAny = z.lazy(() =>
+            z.object({ order: OrderSchema })
+        );
+
+        @Identifier("MutualOrder")
+        @Schema(OrderSchema)
+        class MutualOrder {}
+
+        @Identifier("MutualAddress")
+        @Schema(AddressSchema)
+        class MutualAddress {}
+
+        const orderDef = dataTypeMap(MutualOrder);
+        const addressDef = dataTypeMap(MutualAddress);
+
+        expect(orderDef.type).toBe("{ addresses: MutualAddress[]; }");
+        expect(addressDef.type).toBe("{ order: MutualOrder; }");
+    });
+
+    it("references itself by identifier for self-recursive schemas", () => {
+        const NodeSchema: ZodTypeAny = z.lazy(() =>
+            z.object({ children: z.array(NodeSchema) })
+        );
+
+        @Identifier("TreeNode")
+        @Schema(NodeSchema)
+        class TreeNode {}
+
+        expect(dataTypeMap(TreeNode).type).toBe("{ children: TreeNode[]; }");
+    });
+
+    it("prints registered schemas embedded in other types as identifier references", () => {
+        const ItemSchema: ZodTypeAny = z.lazy(() =>
+            z.object({ related: z.array(ItemSchema) })
+        );
+
+        @Identifier("EmbeddedItem")
+        @Schema(ItemSchema)
+        class EmbeddedItem {}
+
+        @Identifier("ItemPayload")
+        @Schema(z.object({ item: ItemSchema, count: z.number() }))
+        class ItemPayload {}
+
+        dataTypeMap(EmbeddedItem);
+        const payloadDef = dataTypeMap(ItemPayload);
+
+        expect(payloadDef.type).toBe("{ item: EmbeddedItem; count: number; }");
+    });
+
+    it("throws a helpful error when a recursive schema is not registered as a data type", () => {
+        const LoopSchema: ZodTypeAny = z.lazy(() =>
+            z.object({ next: LoopSchema })
+        );
+
+        @Identifier("BrokenPayload")
+        @Schema(z.object({ loop: LoopSchema }))
+        class BrokenPayload {}
+
+        const def = dataTypeMap(BrokenPayload);
+        expect(() => def.type).toThrow(/BrokenPayload.*recursive/i);
+    });
+});