feat: add support recursive types

README.md

@@ -95,7 +95,6 @@ export const api = {
 ## TODO
 
 - ZodNever
-- ZodLazy (recursive types)
 
 ## License
 

lib/MockGenerator.ts

@@ -32,10 +32,13 @@ import PipelineGenerator from './generators/PipelineGenerator';
 import BrandedGenerator from './generators/BrandedGenerator';
 import FunctionGenerator from './generators/FunctionGenerator';
 import LazyGenerator from './generators/LazyGenerator';
+import { DepthLimitError } from './errors/DepthLimitError';
 
+const _schemasCache = new WeakMap<z.ZodTypeAny, any>();
 export default class MockGenerator<T extends z.ZodTypeAny> {
   private generator: BaseGenerator<TypeOf<T>>;
   private schema: T;
+  private readonly MAX_DEPTH = 3;
 
   constructor(schema: T) {
     this.schema = schema;
@@ -86,6 +89,32 @@ export default class MockGenerator<T extends z.ZodTypeAny> {
   }
 
   public generate(): z.infer<T> {
-    return this.generator.generate(this.schema);
+    this.incrementRecursionCount();
+
+    try {
+      const generated = this.generator.generate(this.schema);
+      return generated;
+    }
+    finally {
+      this.decrementRecursionCount();
+    }
+  }
+
+  private incrementRecursionCount(): void {
+    const recursionCount = _schemasCache.get(this.schema) ?? 0;
+    if (recursionCount > this.MAX_DEPTH) {
+      throw new DepthLimitError();
+    }
+
+    _schemasCache.set(this.schema, recursionCount + 1);
+  }
+
+  private decrementRecursionCount(): void {
+    const recursionCount = _schemasCache.get(this.schema) ?? 0;
+    _schemasCache.set(this.schema, recursionCount - 1);
+  }
+
+  public reset(): void {
+    _schemasCache.delete(this.schema);
   }
 }

lib/errors/DepthLimitError.ts

@@ -0,0 +1,6 @@
+export class DepthLimitError extends Error {
+  constructor() {
+    super('Depth limit reached');
+    this.name = 'DepthLimitError';
+  }
+}

lib/generators/ArrayGenerator.ts

@@ -1,5 +1,6 @@
 import type { z } from 'zod';
 import MockGenerator from '../MockGenerator';
+import { DepthLimitError } from '../errors/DepthLimitError';
 import type BaseGenerator from './BaseGenerator';
 
 const DEFAULT_LENGTH = 3;
@@ -16,7 +17,15 @@ export default class ArrayGenerator<T extends z.ZodArray<any>> implements BaseGe
       length = schema._def.exactLength.value;
     }
 
-    const mockGenerator = new MockGenerator(schema._def.type);
-    return Array.from({ length }, () => mockGenerator.generate());
+    try {
+      const mockGenerator = new MockGenerator(schema.element);
+      return Array.from({ length }, () => mockGenerator.generate());
+    }
+    catch (e) {
+      if (e instanceof DepthLimitError) {
+        return [];
+      }
+      throw e;
+    }
   }
 }

lib/generators/ObjectGenerator.ts

@@ -1,13 +1,24 @@
 import type { ZodRawShape, z } from 'zod';
 import MockGenerator from '../MockGenerator';
+import { DepthLimitError } from '../errors/DepthLimitError';
 import type BaseGenerator from './BaseGenerator';
 
 export default class ObjectGenerator<T extends ZodRawShape, U extends z.ZodObject<T>> implements BaseGenerator<U> {
   public generate(schema: U) {
-    return Object.entries(schema._def.shape()).reduce((acc, [key, value]) => {
-      const mockGenerator = new MockGenerator(value);
-      acc[key as keyof z.infer<U>] = mockGenerator.generate();
-      return acc;
-    }, {} as z.infer<U>);
+    const generated: z.infer<U> = {} as z.infer<U>;
+    Object.entries(schema._def.shape()).forEach(([key, value]) => {
+      try {
+        const mockGenerator = new MockGenerator(value);
+        generated[key as keyof z.infer<U>] = mockGenerator.generate();
+      }
+      catch (e) {
+        if (e instanceof DepthLimitError) {
+          return;
+        }
+        throw e;
+      }
+    });
+
+    return generated;
   }
 }

tests/lazy.test.ts

@@ -8,4 +8,16 @@ describe('Lazy', () => {
 
     expect(schema);
   });
+
+  test('recursive', () => {
+    const baseCategorySchema = z.object({ name: z.string() });
+    type Category = z.infer<typeof baseCategorySchema> & {
+      subcategories: Category[]
+    };
+    const categorySchema: z.ZodType<Category> = baseCategorySchema.extend({
+      subcategories: z.array(z.lazy(() => categorySchema)),
+    });
+
+    expect(categorySchema);
+  });
 });