Fix flattening zod schema - chained calls, enum value, tsConfigFilePath

package.json

@@ -27,7 +27,7 @@
   },
   "license": "MIT",
   "engines": {
-    "node": "18 || 19 || 20"
+    "node": "18 || 19 || 20 || 22"
   },
   "workspaces": [
     "packages/*",

packages/nestjs-trpc/lib/generators/__tests__/__snapshots__/procedure.generator.spec.ts.snap

@@ -0,0 +1,32 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`ProcedureGenerator flattenZodSchema should flatten all chained call expressions 1`] = `
+"z.object({
+          options: z
+            .object({
+              userId: z.string().describe('ID of the current user'),
+              type1: z
+          .enum(['Normal', 'Unknown'])
+          .describe('Type of the item').optional().describe('Type 1 of the item')
+            })
+            .merge({
+              z.object({
+                type2: z
+          .enum(['Normal', 'Unknown'])
+          .describe('Type of the item').optional().describe('Type 2 of the item')
+              })
+            })
+            .describe('Options to find many items'),
+        })"
+`;
+
+exports[`ProcedureGenerator flattenZodSchema should flatten enum to literal value 1`] = `
+"z.object({
+          options: z
+            .object({
+              userId: z.string().describe('ID of the current user'),
+              type: z.literal('Normal').describe('Type of the item')
+            })
+            .describe('Options to find many items'),
+        })"
+`;

packages/nestjs-trpc/lib/generators/__tests__/generator.module.spec.ts

@@ -0,0 +1,23 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { GeneratorModule } from '../generator.module';
+
+describe('GeneratorModule', () => {
+  let generatorModule: GeneratorModule;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      imports: [
+        GeneratorModule.forRoot({
+            rootModuleFilePath: '',
+            tsConfigFilePath: './tsconfig.json',
+        }),
+      ],
+    }).compile();
+
+    generatorModule = module.get<GeneratorModule>(GeneratorModule);
+  });
+
+  it('should be defined', () => {
+    expect(generatorModule).toBeDefined();
+  });
+});

packages/nestjs-trpc/lib/generators/__tests__/procedure.generator.spec.ts

@@ -1,5 +1,5 @@
 import { Test, TestingModule } from '@nestjs/testing';
-import { Project } from 'ts-morph';
+import { Identifier, Project, SourceFile, SyntaxKind } from 'ts-morph';
 import {
   ProcedureGeneratorMetadata,
 } from '../../interfaces/generator.interface';
@@ -10,15 +10,14 @@ import { TYPESCRIPT_APP_ROUTER_SOURCE_FILE } from '../generator.constants';
 
 describe('ProcedureGenerator', () => {
   let procedureGenerator: ProcedureGenerator;
-  let project: Project;
 
   beforeEach(async () => {
     const module: TestingModule = await Test.createTestingModule({
       providers: [
         ProcedureGenerator,
         {
           provide: ImportsScanner,
-          useValue: jest.fn(),
+          useValue: new ImportsScanner(), 
         },
         {
           provide: StaticGenerator,
@@ -69,4 +68,75 @@ describe('ProcedureGenerator', () => {
       });
     })
   });
+
+  describe('flattenZodSchema', () => {
+    let project: Project;
+
+    beforeEach(async () => {
+      project = new Project();
+    });
+
+    it('should flatten all chained call expressions', () => {
+      const sourceFile: SourceFile = project.createSourceFile(
+        "test.ts",
+        `
+        import { z } from 'zod';
+
+        const TypeEnum = z
+          .enum(['Normal', 'Unknown'])
+          .describe('Type of the item');
+
+        const FindManyInput = z.object({
+          options: z
+            .object({
+              userId: z.string().describe('ID of the current user'),
+              type1: TypeEnum.optional().describe('Type 1 of the item')
+            })
+            .merge({
+              z.object({
+                type2: TypeEnum.optional().describe('Type 2 of the item')
+              })
+            })
+            .describe('Options to find many items'),
+        });
+        `,
+        { overwrite: true }
+      );
+
+      const node = sourceFile.getDescendantsOfKind(SyntaxKind.Identifier).find((identifier) => identifier.getText() === "FindManyInput") as Identifier;
+      const result = procedureGenerator.flattenZodSchema(node, sourceFile, project, node.getText());
+      expect(result).toMatchSnapshot();
+    });
+
+    it('should flatten enum to literal value', () => {
+      project.createSourceFile(
+        "types.ts",
+        `
+        export enum TypeEnum { Normal = 'Normal', Unknown = 'Unknown' };
+        `,
+        { overwrite: true }
+      );
+      const sourceFile: SourceFile = project.createSourceFile(
+        "test.ts",
+        `
+        import { z } from 'zod';
+        import { TypeEnum } from './types';
+
+        const FindManyInput = z.object({
+          options: z
+            .object({
+              userId: z.string().describe('ID of the current user'),
+              type: z.literal(TypeEnum.Normal).describe('Type of the item')
+            })
+            .describe('Options to find many items'),
+        });
+        `,
+        { overwrite: true }
+      );
+
+      const node = sourceFile.getDescendantsOfKind(SyntaxKind.Identifier).find((identifier) => identifier.getText() === "FindManyInput") as Identifier;
+      const result = procedureGenerator.flattenZodSchema(node, sourceFile, project, node.getText());
+      expect(result).toMatchSnapshot();
+    });
+  });
 });

packages/nestjs-trpc/lib/generators/generator.interface.ts

@@ -6,4 +6,5 @@ export interface GeneratorModuleOptions {
   context?: Class<TRPCContext>;
   outputDirPath?: string;
   schemaFileImports?: Array<SchemaImports>;
+  tsConfigFilePath?: string;
 }

packages/nestjs-trpc/lib/generators/generator.module.ts

@@ -58,7 +58,10 @@ export class GeneratorModule implements OnModuleInit {
       checkJs: true,
       esModuleInterop: true,
     };
-    const project = new Project({ compilerOptions: defaultCompilerOptions });
+    const project = new Project({
+      compilerOptions: defaultCompilerOptions,
+      tsConfigFilePath: options.tsConfigFilePath,
+    });
 
     const appRouterSourceFile = project.createSourceFile(
       path.resolve(options.outputDirPath ?? './', 'server.ts'),

packages/nestjs-trpc/lib/generators/procedure.generator.ts

@@ -1,7 +1,7 @@
 import { Inject, Injectable } from '@nestjs/common';
 import { ProcedureGeneratorMetadata } from '../interfaces/generator.interface';
 import { ProcedureType } from '../trpc.enum';
-import { Project, SourceFile, Node } from 'ts-morph';
+import { Project, SourceFile, Node, SyntaxKind } from 'ts-morph';
 import { ImportsScanner } from '../scanners/imports.scanner';
 import { StaticGenerator } from './static.generator';
 import { TYPESCRIPT_APP_ROUTER_SOURCE_FILE } from './generator.constants';
@@ -48,7 +48,34 @@ export class ProcedureGenerator {
       sourceFile,
       project,
     );
-    if (Node.isIdentifier(node)) {
+    if (Node.isPropertyAccessExpression(node)) {
+      const propertyAccess = node.asKindOrThrow(
+        SyntaxKind.PropertyAccessExpression,
+      );
+      const enumName = propertyAccess.getExpression().getText();
+      const enumDeclaration =
+        sourceFile.getEnum(enumName) ??
+        importsMap
+          .get(enumName)
+          ?.initializer?.asKind(SyntaxKind.EnumDeclaration);
+
+      let enumValue: string | undefined;
+      if (enumDeclaration) {
+        enumValue = enumDeclaration
+          .getMember(propertyAccess.getName())
+          ?.getInitializer()
+          ?.getText();
+      }
+
+      schema =
+        enumValue ??
+        `${this.flattenZodSchema(
+          node.getExpression(),
+          sourceFile,
+          project,
+          node.getExpression().getText(),
+        )}.${propertyAccess.getName()}`;
+    } else if (Node.isIdentifier(node)) {
       const identifierName = node.getText();
       const identifierDeclaration =
         sourceFile.getVariableDeclaration(identifierName);
@@ -114,16 +141,14 @@ export class ProcedureGenerator {
         Node.isPropertyAccessExpression(expression) &&
         !expression.getText().startsWith('z')
       ) {
-        const baseSchema = this.flattenZodSchema(
-          expression,
-          sourceFile,
-          project,
-          expression.getText(),
-        );
-        const propertyName = expression.getName();
         schema = schema.replace(
           expression.getText(),
-          `${baseSchema}.${propertyName}`,
+          this.flattenZodSchema(
+            expression,
+            sourceFile,
+            project,
+            expression.getText(),
+          ),
         );
       } else if (!expression.getText().startsWith('z')) {
         this.staticGenerator.addSchemaImports(
@@ -140,13 +165,20 @@ export class ProcedureGenerator {
           this.flattenZodSchema(arg, sourceFile, project, argText),
         );
       }
-    } else if (Node.isPropertyAccessExpression(node)) {
-      schema = this.flattenZodSchema(
-        node.getExpression(),
-        sourceFile,
-        project,
-        node.getExpression().getText(),
-      );
+
+      for (const child of expression.getChildren()) {
+        if (Node.isCallExpression(child)) {
+          const childText = child.getText();
+          schema = schema.replace(
+            childText,
+            this.flattenZodSchema(child, sourceFile, project, childText),
+          );
+        }
+      }
+    } else if (Node.isFunctionExpression(node)) {
+      schema = 'function() { /* Function body removed */ }';
+    } else if (Node.isArrowFunction(node)) {
+      schema = '() => /* Function body removed */ undefined';
     }
 
     return schema;

packages/nestjs-trpc/lib/interfaces/module-options.interface.ts

@@ -19,6 +19,16 @@ export interface TRPCModuleOptions {
    */
   autoSchemaFile?: string;
 
+  /**
+   * Path to project tsconfig.
+   */
+  tsConfigFilePath?: string;
+
+  /**
+   * The file path of the module that calls the TRPCModule.
+   */
+  rootModuleFilePath?: string;
+
   /**
    * Specifies additional imports for the schema file. This array can include functions, objects, or Zod schemas.
    * While `nestjs-trpc` typically handles imports automatically, this option allows manual inclusion of imports for exceptional cases.

packages/nestjs-trpc/lib/trpc.module.ts

@@ -54,7 +54,8 @@ export class TRPCModule implements OnModuleInit {
       imports.push(
         GeneratorModule.forRoot({
           outputDirPath: options.autoSchemaFile,
-          rootModuleFilePath: callerFilePath,
+          tsConfigFilePath: options.tsConfigFilePath,
+          rootModuleFilePath: options.rootModuleFilePath ?? callerFilePath,
           schemaFileImports: options.schemaFileImports,
           context: options.context,
         }),