fix: import z constant in generated files

stainless-api · Jul 14, 2023 · 49a50a4 · 49a50a4
1 parent aaed9a3
commit 49a50a4
Show file tree

Hide file tree

Showing 12 changed files with 234 additions and 189 deletions.
diff --git a/packages/ts-to-zod/src/__tests__/crazy-prisma-types.test.ts b/packages/ts-to-zod/src/__tests__/crazy-prisma-types.test.ts
diff --git a/packages/ts-to-zod/src/__tests__/generate-in-node-modules.test.ts b/packages/ts-to-zod/src/__tests__/generate-in-node-modules.test.ts
@@ -22,7 +22,7 @@ it(`generating in node_modules`, async () =>
       },
       rootPath: await (async () => {
         const rootPackageJson = await pkgUp({
-            cwd: __dirname
+          cwd: __dirname
         });
         if (!rootPackageJson) {
           throw new Error("test must run within npm package");
@@ -33,11 +33,13 @@ it(`generating in node_modules`, async () =>
   })
 ).toMatchInlineSnapshot(`
 {
-  "node_modules/zodgen/src/__tests__/common.ts": "import { Enum as __enum_Enum } from "../../../../src/__tests__/common.ts";
+  "node_modules/zodgen/src/__tests__/common.ts": "import { z } from "zod";
+import { Enum as __enum_Enum } from "../../../../src/__tests__/common.ts";
 export const AddressIface = z.object({ street: z.string(), city: z.string(), state: z.string(), postalCode: z.string() });
 export const Enum = z.nativeEnum(__enum_Enum);
 ",
-  "node_modules/zodgen/src/__tests__/generate-in-node-modules.test.ts": "import { AddressIface as AddrIface, Enum } from "./common.ts";
+  "node_modules/zodgen/src/__tests__/generate-in-node-modules.test.ts": "import { z } from "zod";
+import { AddressIface as AddrIface, Enum } from "./common.ts";
 const T = z.object({ datetime: z.string().max(20).datetime(), addresses: z.array(z.lazy(() => AddrIface)), enum: z.lazy(() => Enum) });
 ",
 }

diff --git a/packages/ts-to-zod/src/__tests__/imported-aliased-interface.test.ts b/packages/ts-to-zod/src/__tests__/imported-aliased-interface.test.ts
@@ -18,15 +18,17 @@ interface T {
 
 it(`imported interface`, async () =>
   expect(
-    await multiFileTestCase({
-      __filename,
-    })
-  ).toMatchInlineSnapshot(`
-    {
-      "src/__tests__/common.codegen.ts": "export const AddressIface = z.object({ street: z.string(), city: z.string(), state: z.string(), postalCode: z.string() });
-    ",
-      "src/__tests__/imported-aliased-interface.test.codegen.ts": "import { AddressIface as Addr } from "./common.codegen.ts";
-    const T = z.object({ firstName: z.string(), lastName: z.string(), address: z.lazy(() => Addr).optional() });
-    ",
-    }
-  `));
+  await multiFileTestCase({
+    __filename
+  })
+).toMatchInlineSnapshot(`
+{
+  "src/__tests__/common.codegen.ts": "import { z } from "zod";
+export const AddressIface = z.object({ street: z.string(), city: z.string(), state: z.string(), postalCode: z.string() });
+",
+  "src/__tests__/imported-aliased-interface.test.codegen.ts": "import { z } from "zod";
+import { AddressIface as Addr } from "./common.codegen.ts";
+const T = z.object({ firstName: z.string(), lastName: z.string(), address: z.lazy(() => Addr).optional() });
+",
+}
+`));
diff --git a/packages/ts-to-zod/src/__tests__/imported-enum.test.ts b/packages/ts-to-zod/src/__tests__/imported-enum.test.ts
@@ -5,16 +5,18 @@ type T = Enum[];
 
 it(`imported enum`, async () =>
   expect(
-    await multiFileTestCase({
-      __filename,
-    })
-  ).toMatchInlineSnapshot(`
-    {
-      "src/__tests__/common.codegen.ts": "import { Enum as __enum_Enum } from "./common.ts";
-    export const Enum = z.nativeEnum(__enum_Enum);
-    ",
-      "src/__tests__/imported-enum.test.codegen.ts": "import { Enum } from "./common.codegen.ts";
-    const T = z.array(z.lazy(() => Enum));
-    ",
-    }
-  `));
+  await multiFileTestCase({
+    __filename
+  })
+).toMatchInlineSnapshot(`
+{
+  "src/__tests__/common.codegen.ts": "import { z } from "zod";
+import { Enum as __enum_Enum } from "./common.ts";
+export const Enum = z.nativeEnum(__enum_Enum);
+",
+  "src/__tests__/imported-enum.test.codegen.ts": "import { z } from "zod";
+import { Enum } from "./common.codegen.ts";
+const T = z.array(z.lazy(() => Enum));
+",
+}
+`));
diff --git a/packages/ts-to-zod/src/__tests__/imported-interface.test.ts b/packages/ts-to-zod/src/__tests__/imported-interface.test.ts
@@ -25,9 +25,11 @@ it(`imported interface`, async () =>
   })
 ).toMatchInlineSnapshot(`
 {
-  "src/__tests__/common.codegen.ts": "export const AddressIface = z.object({ street: z.string(), city: z.string(), state: z.string(), postalCode: z.string() });
+  "src/__tests__/common.codegen.ts": "import { z } from "zod";
+export const AddressIface = z.object({ street: z.string(), city: z.string(), state: z.string(), postalCode: z.string() });
 ",
-  "src/__tests__/imported-interface.test.codegen.ts": "import { AddressIface } from "./common.codegen.ts";
+  "src/__tests__/imported-interface.test.codegen.ts": "import { z } from "zod";
+import { AddressIface } from "./common.codegen.ts";
 import { TestClass as __class_TestClass } from "./imported-interface.test.ts";
 const T = z.object({ firstName: z.string(), lastName: z.string(), address: z.lazy(() => AddressIface).optional(), testMyClass: z.instanceof(__class_TestClass) });
 ",

diff --git a/packages/ts-to-zod/src/__tests__/imported-type.test.ts b/packages/ts-to-zod/src/__tests__/imported-type.test.ts
@@ -9,15 +9,17 @@ type T = {
 
 it(`imported type`, async () =>
   expect(
-    await multiFileTestCase({
-      __filename,
-    })
-  ).toMatchInlineSnapshot(`
-    {
-      "src/__tests__/common.codegen.ts": "export const Address = z.object({ street: z.string(), city: z.string(), state: z.string(), postalCode: z.string() });
-    ",
-      "src/__tests__/imported-type.test.codegen.ts": "import { Address } from "./common.codegen.ts";
-    const T = z.object({ firstName: z.string(), lastName: z.string(), address: z.lazy(() => Address).optional() });
-    ",
-    }
-  `));
+  await multiFileTestCase({
+    __filename
+  })
+).toMatchInlineSnapshot(`
+{
+  "src/__tests__/common.codegen.ts": "import { z } from "zod";
+export const Address = z.object({ street: z.string(), city: z.string(), state: z.string(), postalCode: z.string() });
+",
+  "src/__tests__/imported-type.test.codegen.ts": "import { z } from "zod";
+import { Address } from "./common.codegen.ts";
+const T = z.object({ firstName: z.string(), lastName: z.string(), address: z.lazy(() => Address).optional() });
+",
+}
+`));
diff --git a/packages/ts-to-zod/src/__tests__/mutual-recursion.test.ts b/packages/ts-to-zod/src/__tests__/mutual-recursion.test.ts
@@ -5,13 +5,14 @@ type U = string | [T];
 
 it(`mutual recursion`, async () =>
   expect(
-    await multiFileTestCase({
-      __filename,
-    })
-  ).toMatchInlineSnapshot(`
-    {
-      "src/__tests__/mutual-recursion.test.codegen.ts": "const U = z.union([z.string(), z.tuple([z.lazy(() => T)])]);
-    const T = z.union([z.number(), z.tuple([z.lazy(() => U)])]);
-    ",
-    }
-  `));
+  await multiFileTestCase({
+    __filename
+  })
+).toMatchInlineSnapshot(`
+{
+  "src/__tests__/mutual-recursion.test.codegen.ts": "import { z } from "zod";
+const U = z.union([z.string(), z.tuple([z.lazy(() => T)])]);
+const T = z.union([z.number(), z.tuple([z.lazy(() => U)])]);
+",
+}
+`));
diff --git a/packages/ts-to-zod/src/__tests__/native-types.test.ts b/packages/ts-to-zod/src/__tests__/native-types.test.ts
@@ -7,12 +7,13 @@ type T = {
 
 it(`native types`, async () =>
   expect(
-    await multiFileTestCase({
-      __filename,
-    })
-  ).toMatchInlineSnapshot(`
-    {
-      "src/__tests__/native-types.test.codegen.ts": "const T = z.object({ map: z.map(z.string(), z.number()), date: z.date() });
-    ",
-    }
-  `));
+  await multiFileTestCase({
+    __filename
+  })
+).toMatchInlineSnapshot(`
+{
+  "src/__tests__/native-types.test.codegen.ts": "import { z } from "zod";
+const T = z.object({ map: z.map(z.string(), z.number()), date: z.date() });
+",
+}
+`));
diff --git a/packages/ts-to-zod/src/__tests__/transform-refine.test.ts b/packages/ts-to-zod/src/__tests__/transform-refine.test.ts
@@ -78,7 +78,8 @@ it(`transform`, async () =>
   })
 ).toMatchInlineSnapshot(`
 {
-  "src/__tests__/transform-refine.test.codegen.ts": "import { ParseFloat, ToString, Coerce, ParsePet, Even } from "./transform-refine.test.ts";
+  "src/__tests__/transform-refine.test.codegen.ts": "import { z } from "zod";
+import { ParseFloat, ToString, Coerce, ParsePet, Even } from "./transform-refine.test.ts";
 const T = z.object({ a: z.date().transform(new ToString().transform).transform(new ParseFloat().transform), b: z.string().transform(new Coerce().transform), c: z.string().refine(new ParsePet().refine, new ParsePet().message), d: z.number().superRefine(new Even().superRefine), date: z.date().min(new Date("2023-01-10")), number: z.number().finite().safe("too big").gt(5, "5 and below too small!"), object: z.object({}).passthrough(), array: z.array(z.number().nullable()).nonempty().min(5, "at least five elements needed"), datetime: z.string().datetime({ offset: true }), catchall: z.object({ a: z.number() }).catchall(z.string()) });
 ",
 }

diff --git a/packages/ts-to-zod/src/__tests__/xor.test.ts b/packages/ts-to-zod/src/__tests__/xor.test.ts
@@ -28,17 +28,18 @@ type T = XOR<Person, Pet>;
 
 it(`XOR type`, async () => {
   expect(
-    await multiFileTestCase({
-      __filename,
-    })
-  ).toMatchInlineSnapshot(`
-    {
-      "src/__tests__/xor.test.codegen.ts": "const Pet = z.object({ name: z.string(), breed: z.string() });
-    const Person = z.object({ name: z.string(), language: z.string() });
-    const T = z.union([z.object({ language: z.undefined() }).and(z.lazy(() => Pet)), z.object({ breed: z.undefined() }).and(z.lazy(() => Person))]);
-    ",
-    }
-  `);
+  await multiFileTestCase({
+    __filename
+  })
+).toMatchInlineSnapshot(`
+{
+  "src/__tests__/xor.test.codegen.ts": "import { z } from "zod";
+const Pet = z.object({ name: z.string(), breed: z.string() });
+const Person = z.object({ name: z.string(), language: z.string() });
+const T = z.union([z.object({ language: z.undefined() }).and(z.lazy(() => Pet)), z.object({ breed: z.undefined() }).and(z.lazy(() => Person))]);
+",
+}
+`);
 
   const Pet = z.object({ name: z.string(), breed: z.string() });
   const Person = z.object({ name: z.string(), language: z.string() });

diff --git a/packages/ts-to-zod/src/filePathConfig.ts b/packages/ts-to-zod/src/filePathConfig.ts
@@ -14,6 +14,8 @@ export interface GenOptions {
    * paths are resolved
    */
   rootPath: string;
+  /** the package from where to import the `z` constant */
+  zPackage?: string;
 }
 
 // produced via the _

diff --git a/packages/ts-to-zod/src/generateFiles.ts b/packages/ts-to-zod/src/generateFiles.ts
@@ -3,8 +3,11 @@ import { ImportInfo, SchemaGenContext } from "./convertType";
 import { ts, Symbol } from "ts-morph";
 const { factory } = ts;
 import * as Path from "path";
-import { GenOptions, GenerationConfig, createGenerationConfig } from "./filePathConfig";
-
+import {
+  GenOptions,
+  GenerationConfig,
+  createGenerationConfig,
+} from "./filePathConfig";
 
 export function generateFiles(
   ctx: SchemaGenContext,
@@ -15,10 +18,15 @@ export function generateFiles(
   for (const [path, info] of ctx.files.entries()) {
     const generatedPath = generatePath({
       path,
-      ...generationConfig
+      ...generationConfig,
     });
 
-    const statements: ts.Statement[] = generateImportStatements(generationConfig, generatedPath, info.imports);
+    const statements: ts.Statement[] = generateImportStatements(
+      generationConfig,
+      generatedPath,
+      options.zPackage,
+      info.imports
+    );
 
     for (const schema of info.generatedSchemas) {
       const declaration = factory.createVariableDeclaration(
@@ -93,6 +101,7 @@ function generatePath({
 export function generateImportStatements(
   config: GenerationConfig,
   filePath: string,
+  zPackage: string | undefined,
   imports: Map<Symbol, ImportInfo>
 ): ts.ImportDeclaration[] {
   const importDeclarations = [];
@@ -105,10 +114,28 @@ export function generateImportStatements(
           ? symbol.getDeclarations()[0].getSourceFile().getFilePath()
           : generatePath({
               path: symbol.getDeclarations()[0].getSourceFile().getFilePath(),
-              ...config
+              ...config,
             })
       )
   );
+  const zImportClause = factory.createImportClause(
+    false,
+    undefined,
+    factory.createNamedImports([
+      factory.createImportSpecifier(
+        false,
+        undefined,
+        factory.createIdentifier("z")
+      ),
+    ])
+  );
+  const zImportDeclaration = factory.createImportDeclaration(
+    [],
+    zImportClause,
+    factory.createStringLiteral(zPackage || "zod")
+  );
+
+  importDeclarations.push(zImportDeclaration);
 
   for (const [relativePath, entries] of Object.entries(importGroups)) {
     const importSpecifiers = entries.map(([symbol, { as }]) => {