Repository: fabien0102/ts-to-zod
Pull request #148 (2023-07-24 20:38:13)

Handling imports of generated types
# Why

Following #135 (once it's merged), this PR adds support for automatic resolution of ts-to-zod-generated schemas, based on the config file provided: if the current input file imports a type from a file referenced in the `ts-to-zod.config.js` file, then the generated file will automatically resolve the reference to the corresponding schema. 

```ts

//ts-to-zod.config.js
/**
 * ts-to-zod configuration.
 *
 * @type {import("./src/config").TsToZodConfig}
 */
module.exports = [
  {
    name: "villain",
    input: "src/villain.ts",
    output: "src/generated/villain.zod.ts",
    getSchemaName: (id) => `z${id}`
  },
   {
    name: "hero",
    input: "src/example/heros.ts",
    output: "src/example/heros.zod.ts",
  },
];

// heros.ts (input)
import { Villain } from "../villain";

export interface Hero {
  name: string;
  nemesis: Villain;
}

// heros.zod.ts (output)
import { zVillain } from "../generated/villain.zod";

export const heroSchema = z.object({
  name: z.string(),
  nemesis: zVillain;
});
``` 

This opens up the possibility of handling multiple files in one generation (by adding them to the config file first for instance). 


tvillaren: Hello @dan-blank, @fabien0102 👋, 

Following the discussion in #135, I fiddled a bit and decided to go for this solution of Zod Schemas. 
Happy to discuss this.
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/fabien0102/ts-to-zod/pull/148?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD) Report
> Merging [#148](https://app.codecov.io/gh/fabien0102/ts-to-zod/pull/148?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD) (43e6762) into [main](https://app.codecov.io/gh/fabien0102/ts-to-zod/commit/8d3b297e4131ef105bfe155b038e4b4f76a7d5fc?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD) (8d3b297) will **increase** coverage by `0.06%`.
> The diff coverage is `97.41%`.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD).

```diff
@@            Coverage Diff             @@
##             main     #148      +/-   ##
==========================================
+ Coverage   97.33%   97.40%   +0.06%     
==========================================
  Files          14       15       +1     
  Lines         676      770      +94     
  Branches      275      310      +35     
==========================================
+ Hits          658      750      +92     
- Misses         18       20       +2     
```


| [Files Changed](https://app.codecov.io/gh/fabien0102/ts-to-zod/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD) | Coverage Δ | |
|---|---|---|
| [src/core/generateZodSchema.ts](https://app.codecov.io/gh/fabien0102/ts-to-zod/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD#diff-c3JjL2NvcmUvZ2VuZXJhdGVab2RTY2hlbWEudHM=) | `95.88% <75.00%> (-0.32%)` | :arrow_down: |
| [src/utils/fixOptional.ts](https://app.codecov.io/gh/fabien0102/ts-to-zod/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD#diff-c3JjL3V0aWxzL2ZpeE9wdGlvbmFsLnRz) | `95.00% <95.00%> (ø)` | |
| [src/core/generate.ts](https://app.codecov.io/gh/fabien0102/ts-to-zod/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD#diff-c3JjL2NvcmUvZ2VuZXJhdGUudHM=) | `97.20% <100.00%> (+1.08%)` | :arrow_up: |
| [src/core/validateGeneratedTypes.ts](https://app.codecov.io/gh/fabien0102/ts-to-zod/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD#diff-c3JjL2NvcmUvdmFsaWRhdGVHZW5lcmF0ZWRUeXBlcy50cw==) | `97.61% <100.00%> (ø)` | |
| [src/utils/importHandling.ts](https://app.codecov.io/gh/fabien0102/ts-to-zod/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Fabien+BERNARD#diff-c3JjL3V0aWxzL2ltcG9ydEhhbmRsaW5nLnRz) | `100.00% <100.00%> (ø)` | |


anthony-dandrea: This worked pretty well for me when testing in my project. One issue I ran into is that I use [proto-loader-gen-types](https://github.com/grpc/grpc-node/blob/master/packages/proto-loader/README.md) for generating some of my types from protos. I found that it doesn't work with there are parenthesis around the types which is how proto-loader-gen-types generates types. Example:

```js
// DoesntWork.ts
import type { BoolValue } from "google/protobuf/BoolValue";
export interface Broken {
  bool: (BoolValue | null);
}

// DoesWork.ts
import type { BoolValue } from "google/protobuf/BoolValue";
export interface Works {
  bool: BoolValue | null;
}

// google/protobuf/BoolValue.ts
export interface BoolValue {
  'value': (boolean);
}
```
tvillaren: Hey @anthony-dandrea, 

The issue you mentioned was not directly related to this PR, but it's an easy fix. 
The fix is there: https://github.com/fabien0102/ts-to-zod/pull/152

You can cherry pick the commit in your branch while waiting for an approval. I'll rebase the import-related branches when it's  the fix is merged. 
schiller-manuel: Can you please explain the difference between this PR and https://github.com/fabien0102/ts-to-zod/pull/164 ?
Would love to get all those PRs merged.
tvillaren: Hey @schiller-manuel,

As far as I see, it seems that they are multiple PRs to handle import: 

Those 3 are not from me, and they felt like too much changes in one pass to me. 
- https://github.com/fabien0102/ts-to-zod/pull/118
- https://github.com/fabien0102/ts-to-zod/pull/126
- https://github.com/fabien0102/ts-to-zod/pull/164 based on https://github.com/fabien0102/ts-to-zod/pull/118

I took a different approach with https://github.com/fabien0102/ts-to-zod/pull/135 + https://github.com/fabien0102/ts-to-zod/pull/148 (the current one):

👉 https://github.com/fabien0102/ts-to-zod/pull/135 aimed at handling `import` statement *in one file* only. It does work, as I've been using this for months now on my company project. 
However, due to validation errors that I haven't been able to fix (see https://github.com/fabien0102/ts-to-zod/issues/186), it would be kind of a breaking change in the current state, as non-zod imported types would only pass validation in most cases (the generated zod validators usually work in that case).

👉 the current PR https://github.com/fabien0102/ts-to-zod/pull/148 adds a very simple way of handling multiple file generation, by relying on the `ts-to-zod.config.js` configuration file to describe all the files in the "ZOD filesystem". I did that as a first step, to rely on existing bricks and not do a big bang refactoring like other approaches above. 
Again, this work in my project. I haven't rebased to the latest `ts-to-zod` though because of the validation errors in https://github.com/fabien0102/ts-to-zod/pull/135 

👉 in my mind, a third PR would have finished the work to automate filesystem browsing and generate all zod files in the right order based on `import` statements and all. Haven't looked at it yet (but the need in my company project is going to arise to if we can move on with the rest, it would be great 😉)

## How to move on

I see several alternatives to move on with my 2 PRs: 
1. Find a solution to external imports validation issues, as discussed on https://github.com/fabien0102/ts-to-zod/issues/186, implement in https://github.com/fabien0102/ts-to-zod/pull/135, and then integrate to https://github.com/fabien0102/ts-to-zod/pull/148
2. Drop the "external import" bit and focus only on "internal imports" of files which are also converted to zod schemas. This would mean a bit of rework, it would add support for multiple generated files (through `ts-to-zod.config.js` for now), but limited to referencing other "internal" types.  
3. Acknowledge that we don't fully support external imports and merge both PRs (after rebasing).

Any thoughts welcome are welcome!

tvillaren: Hello @schiller-manuel, @fabien0102, 

I rebased this PR and it's (finally) ready for merging. 
Now `ts-to-zod` fully supports imports of types that are also converted into zod schema on the condition that all files are listed in the configuration file. 

Next step would be to build the list "on the fly" (from a folder passed as argument for instance).

tvillaren: I don't think I will have time in the next few weeks to deep dive on this but would love to at one point!
Let's merge, I know some people (including me) are using this branch in their project. 

Patch:
diff --git a/README.md b/README.md
index 3046b96..e976eb2 100644
--- a/README.md
+++ b/README.md
@@ -370,7 +370,9 @@ To resume, you can use all the primitive types and some the following typescript
 - `Array<>`
 - `Promise<>`
 
-This utils is design to work with one file only, and will reference types from the same file:
+## Type references
+
+This utility is designed to work with one file at a time (it will not split one file into several), and will handle references to types from the same file:
 
 ```ts
 // source.ts
@@ -388,7 +390,9 @@ export const heroSchema = z.object({
 });
 ```
 
-It can also reference imported types from any modules but will use an `any` validation as placeholder, as we cannot now
+### Non-Zod imports
+
+`ts-to-zod` can reference imported types from other modules but will use an `any` validation as placeholder, as we cannot now
 their schema.
 
 ```ts
@@ -415,6 +419,54 @@ export const heroSchema = z.object({
 });
 ```
 
+### Zod Imports
+
+If an imported type is referenced in the `ts-to-zod.config.js` config (as input), this utility will automatically replace the import with the given output from the file, resolving the relative paths between both.
+
+```ts
+
+//ts-to-zod.config.js
+/**
+ * ts-to-zod configuration.
+ *
+ * @type {import("./src/config").TsToZodConfig}
+ */
+module.exports = [
+  {
+    name: "villain",
+    input: "src/villain.ts",
+    output: "src/generated/villain.zod.ts",
+    getSchemaName: (id) => `z${id}`
+  },
+   {
+    name: "hero",
+    input: "src/example/heros.ts",
+    output: "src/example/heros.zod.ts",
+  },
+];
+
+// heros.ts (input)
+import { Person } from "@3rdparty/person";
+import { Villain } from "../villain";
+
+export interface Hero {
+  name: string;
+  realPerson: Person;
+  nemesis: Villain;
+}
+
+// heros.zod.ts (output)
+import { zVillain } from "../generated/villain.zod";
+
+const personSchema = z.any();
+
+export const heroSchema = z.object({
+  name: z.string(),
+  realPerson: personSchema,
+  nemesis: zVillain;
+});
+```
+
 ## Programmatic API
 
 You need more than one file? Want even more power? No problem, just use the tool as a library.
diff --git a/example/heros.ts b/example/heros.ts
index 585a56a..e397bb3 100644
--- a/example/heros.ts
+++ b/example/heros.ts
@@ -1,3 +1,5 @@
+import { Person } from "./person";
+
 export enum EnemyPower {
   Flight = "flight",
   Strength = "strength",
@@ -11,7 +13,7 @@ export namespace Skills {
   };
 }
 
-export interface Enemy {
+export interface Enemy extends Person {
   name: string;
   powers: EnemyPower[];
   inPrison: boolean;
diff --git a/example/heros.zod.ts b/example/heros.zod.ts
index 762ad6b..e1004e6 100644
--- a/example/heros.zod.ts
+++ b/example/heros.zod.ts
@@ -2,13 +2,15 @@
 import { z } from "zod";
 import { EnemyPower, Villain, EvilPlan, EvilPlanDetails } from "./heros";
 
+import { personSchema } from "./person.zod";
+
 export const enemyPowerSchema = z.nativeEnum(EnemyPower);
 
 export const skillsSpeedEnemySchema = z.object({
   power: z.literal(EnemyPower.Speed),
 });
 
-export const enemySchema = z.object({
+export const enemySchema = personSchema.extend({
   name: z.string(),
   powers: z.array(enemyPowerSchema),
   inPrison: z.boolean(),
diff --git a/example/person.ts b/example/person.ts
new file mode 100644
index 0000000..b24544d
--- /dev/null
+++ b/example/person.ts
@@ -0,0 +1,4 @@
+// Used as an import in hero.ts
+export interface Person {
+  realName: string;
+}
diff --git a/example/person.zod.ts b/example/person.zod.ts
new file mode 100644
index 0000000..71715f9
--- /dev/null
+++ b/example/person.zod.ts
@@ -0,0 +1,6 @@
+// Generated by ts-to-zod
+import { z } from "zod";
+
+export const personSchema = z.object({
+  realName: z.string(),
+});
diff --git a/src/cli.ts b/src/cli.ts
index 68ceed4..92c4e18 100644
--- a/src/cli.ts
+++ b/src/cli.ts
@@ -7,7 +7,7 @@ import { join, parse, relative } from "path";
 import prettier from "prettier";
 import slash from "slash";
 import ts from "typescript";
-import { Config, TsToZodConfig } from "./config";
+import { Config, TsToZodConfig, InputOutputMapping } from "./config";
 import {
   getSchemaNameSchema,
   nameFilterSchema,
@@ -133,6 +133,8 @@ class TsToZod extends Command {
 
     const fileConfig = await this.loadFileConfig(config, flags);
 
+    const ioMappings = getInputOutputMappings(config);
+
     if (Array.isArray(fileConfig)) {
       if (args.input || args.output) {
         this.error(`INPUT and OUTPUT arguments are not compatible with --all`);
@@ -141,7 +143,7 @@ class TsToZod extends Command {
         await Promise.all(
           fileConfig.map(async (config) => {
             this.log(`Generating "${config.name}"`);
-            const result = await this.generate(args, config, flags);
+            const result = await this.generate(args, config, flags, ioMappings);
             if (result.success) {
               this.log(` 🎉 Zod schemas generated!`);
             } else {
@@ -156,7 +158,7 @@ class TsToZod extends Command {
         this.error(error);
       }
     } else {
-      const result = await this.generate(args, fileConfig, flags);
+      const result = await this.generate(args, fileConfig, flags, ioMappings);
       if (result.success) {
         this.log(`🎉 Zod schemas generated!`);
       } else {
@@ -177,7 +179,7 @@ class TsToZod extends Command {
           ? fileConfig.find((i) => i.input === slash(path))
           : fileConfig;
 
-        const result = await this.generate(args, config, flags);
+        const result = await this.generate(args, config, flags, ioMappings);
         if (result.success) {
           this.log(`🎉 Zod schemas generated!`);
         } else {
@@ -193,11 +195,13 @@ class TsToZod extends Command {
    * @param args
    * @param fileConfig
    * @param Flags
+   * @param inputOutputMappings
    */
   async generate(
     args: { input?: string; output?: string },
     fileConfig: Config | undefined,
-    Flags: Interfaces.InferredFlags<typeof TsToZod.flags>
+    Flags: Interfaces.InferredFlags<typeof TsToZod.flags>,
+    inputOutputMappings: InputOutputMapping[]
   ): Promise<{ success: true } | { success: false; error: string }> {
     const input = args.input || fileConfig?.input;
     const output = args.output || fileConfig?.output;
@@ -214,6 +218,17 @@ See more help with --help`,
     const inputPath = join(process.cwd(), input);
     const outputPath = join(process.cwd(), output || input);
 
+    const relativeIOMappings = inputOutputMappings.map((io) => {
+      const relativeInput = getImportPath(inputPath, io.input);
+      const relativeOutput = getImportPath(outputPath, io.output);
+
+      return {
+        input: relativeInput,
+        output: relativeOutput,
+        getSchemaName: io.getSchemaName,
+      };
+    });
+
     // Check args/flags file extensions
     const extErrors: { path: string; expectedExtensions: string[] }[] = [];
     if (!hasExtensions(input, typescriptExtensions)) {
@@ -250,6 +265,7 @@ See more help with --help`,
 
     const generateOptions: GenerateProps = {
       sourceText,
+      inputOutputMappings: relativeIOMappings,
       ...fileConfig,
     };
     if (typeof Flags.keepComments === "boolean") {
@@ -443,4 +459,22 @@ function hasExtensions(path: string, extensions: string[]) {
   return extensions.includes(ext);
 }
 
+function getInputOutputMappings(
+  config: TsToZodConfig | undefined
+): InputOutputMapping[] {
+  if (!config) {
+    return [];
+  }
+
+  if (Array.isArray(config)) {
+    return config.map((c) => {
+      const { input, output, getSchemaName } = c;
+      return { input, output, getSchemaName };
+    });
+  }
+
+  const { input, output, getSchemaName } = config as Config;
+  return [{ input, output, getSchemaName }];
+}
+
 export = TsToZod;
diff --git a/src/config.ts b/src/config.ts
index cfe66ee..70ba3be 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -113,4 +113,9 @@ export type Configs = Array<
   }
 >;
 
+export type InputOutputMapping = Pick<
+  Config,
+  "input" | "output" | "getSchemaName"
+>;
+
 export type TsToZodConfig = Config | Configs;
diff --git a/src/config.zod.ts b/src/config.zod.ts
index c6b077f..2a01e6a 100644
--- a/src/config.zod.ts
+++ b/src/config.zod.ts
@@ -54,4 +54,10 @@ export const configsSchema = z.array(
   )
 );
 
+export const inputOutputMappingSchema = configSchema.pick({
+  input: true,
+  output: true,
+  getSchemaName: true,
+});
+
 export const tsToZodConfigSchema = z.union([configSchema, configsSchema]);
diff --git a/src/core/generate.ts b/src/core/generate.ts
index e909edf..83aae7b 100644
--- a/src/core/generate.ts
+++ b/src/core/generate.ts
@@ -1,7 +1,13 @@
 import { camel } from "case";
 import { getJsDoc } from "tsutils";
 import ts from "typescript";
-import { JSDocTagFilter, NameFilter, CustomJSDocFormatTypes } from "../config";
+import { relative } from "path";
+import {
+  InputOutputMapping,
+  JSDocTagFilter,
+  NameFilter,
+  CustomJSDocFormatTypes,
+} from "../config";
 import { getSimplifiedJsDocTags } from "../utils/getSimplifiedJsDocTags";
 import { resolveModules } from "../utils/resolveModules";
 import {
@@ -10,7 +16,11 @@ import {
   TypeNode,
 } from "../utils/traverseTypes";
 
-import { getImportIdentifiers } from "../utils/importHandling";
+import {
+  getImportIdentifiers,
+  createImportNode,
+} from "../utils/importHandling";
+
 import { generateIntegrationTests } from "./generateIntegrationTests";
 import { generateZodInferredType } from "./generateZodInferredType";
 import {
@@ -19,6 +29,8 @@ import {
 } from "./generateZodSchema";
 import { transformRecursiveSchema } from "./transformRecursiveSchema";
 
+const DEFAULT_GET_SCHEMA = (id: string) => camel(id) + "Schema";
+
 export interface GenerateProps {
   /**
    * Content of the typescript source file.
@@ -57,11 +69,16 @@ export interface GenerateProps {
    * Path of z.infer<> types file.
    */
   inferredTypes?: string;
-
   /**
    * Custom JSDoc format types.
    */
   customJSDocFormatTypes?: CustomJSDocFormatTypes;
+
+  /**
+   * Map of input/output from config that can
+   * be used to automatically handle imports
+   */
+  inputOutputMappings?: InputOutputMapping[];
 }
 
 /**
@@ -73,10 +90,11 @@ export function generate({
   sourceText,
   nameFilter = () => true,
   jsDocTagFilter = () => true,
-  getSchemaName = (id) => camel(id) + "Schema",
+  getSchemaName = DEFAULT_GET_SCHEMA,
   keepComments = false,
   skipParseJSDoc = false,
   customJSDocFormatTypes = {},
+  inputOutputMappings = [],
 }: GenerateProps) {
   // Create a source file and deal with modules
   const sourceFile = resolveModules(sourceText);
@@ -89,7 +107,9 @@ export function generate({
 
   // handling existing import statements
   const importNamesAvailable = new Set<string>();
+  const importNodes: ts.ImportDeclaration[] = [];
   const importNamesUsed: string[] = [];
+  const importedZodNamesAvailable = new Map<string, string>();
 
   const typesNeedToBeExtracted = new Set<string>();
 
@@ -98,8 +118,34 @@ export function generate({
       typeNameMapping.set(node.name.text, node);
     }
     if (ts.isImportDeclaration(node) && node.importClause) {
-      const imports = getImportIdentifiers(node);
-      imports.forEach((i) => importNamesAvailable.add(i));
+      // Check if we're importing from a mapped file
+      const eligibleMapping = inputOutputMappings.find(
+        (io: InputOutputMapping) =>
+          relative(
+            io.input,
+            (node.moduleSpecifier as ts.StringLiteral).text
+          ) === ""
+      );
+      if (eligibleMapping) {
+        const schemaMethod =
+          eligibleMapping.getSchemaName || DEFAULT_GET_SCHEMA;
+
+        const identifiers = getImportIdentifiers(node);
+        identifiers.forEach((i) =>
+          importedZodNamesAvailable.set(i, schemaMethod(i))
+        );
+
+        const importNode = createImportNode(
+          identifiers.map(schemaMethod),
+          eligibleMapping.output
+        );
+        importNodes.push(importNode); // We assume all identifiers will be used so pushing the whole node
+      }
+      // Not a Zod import, handling it as 3rd party import later on
+      else {
+        const identifiers = getImportIdentifiers(node);
+        identifiers.forEach((i) => importNamesAvailable.add(i));
+      }
     }
   };
 
@@ -123,18 +169,33 @@ export function generate({
   };
   ts.forEachChild(sourceFile, visitor);
 
+  const importedZodSchemas = new Set<string>();
+
   typesNeedToBeExtracted.forEach((typeName) => {
     const node = typeNameMapping.get(typeName);
     if (node) {
       nodes.push(node);
-    } else {
-      if (importNamesAvailable.has(typeName)) {
-        importNamesUsed.push(typeName);
-      }
+      return;
+    }
+    if (importNamesAvailable.has(typeName)) {
+      importNamesUsed.push(typeName);
+      return;
+    }
+    if (importedZodNamesAvailable.has(typeName)) {
+      importedZodSchemas.add(importedZodNamesAvailable.get(typeName) as string);
+      return;
     }
   });
 
   // Generate zod schemas for type nodes
+
+  const getDependencyName = (identifierName: string) => {
+    if (importedZodNamesAvailable.has(identifierName)) {
+      return importedZodNamesAvailable.get(identifierName) as string;
+    }
+    return getSchemaName(identifierName);
+  };
+
   const zodTypeSchemas = nodes.map((node) => {
     const typeName = node.name.text;
     const varName = getSchemaName(typeName);
@@ -143,7 +204,7 @@ export function generate({
       node,
       sourceFile,
       varName,
-      getDependencyName: getSchemaName,
+      getDependencyName: getDependencyName,
       skipParseJSDoc,
       customJSDocFormatTypes,
     });
@@ -151,7 +212,7 @@ export function generate({
     return { typeName, varName, ...zodSchema };
   });
 
-  // Generate zod schemas for imports
+  // Generate zod schemas for 3rd party imports
   const zodImportSchemas = importNamesUsed.map((importName) => {
     const varName = getSchemaName(importName);
     return {
@@ -199,7 +260,8 @@ export function generate({
           const isCircular = dependencies.includes(varName);
           const notGeneratedDependencies = dependencies
             .filter((dep) => dep !== varName)
-            .filter((dep) => !statements.has(dep));
+            .filter((dep) => !statements.has(dep))
+            .filter((dep) => !importedZodSchemas.has(dep));
           if (notGeneratedDependencies.length === 0) {
             done = false;
             if (isCircular) {
@@ -279,9 +341,13 @@ ${
     ? `import { ${typeImportsValues.join(", ")} } from "${typesImportPath}";\n`
     : ""
 }
-${Array.from(statements.values())
-  .map((statement) => print(statement.value))
-  .join("\n\n")}
+${
+  importNodes.length
+    ? importNodes.map((node) => print(node)).join("\n") + "\n\n"
+    : ""
+}${Array.from(statements.values())
+    .map((statement) => print(statement.value))
+    .join("\n\n")}
 `;
 
   const testCases = generateIntegrationTests(
diff --git a/src/core/validateGeneratedTypes.ts b/src/core/validateGeneratedTypes.ts
index a1ba66a..c297052 100644
--- a/src/core/validateGeneratedTypes.ts
+++ b/src/core/validateGeneratedTypes.ts
@@ -7,6 +7,7 @@ import ts from "typescript";
 import { join, sep, posix } from "path";
 import { resolveDefaultProperties } from "../utils/resolveDefaultProperties";
 import { fixOptionalAny } from "../utils/fixOptionalAny";
+
 interface File {
   sourceText: string;
   relativePath: string;
diff --git a/src/utils/importHandling.ts b/src/utils/importHandling.ts
index 132f0e4..6a25484 100644
--- a/src/utils/importHandling.ts
+++ b/src/utils/importHandling.ts
@@ -1,4 +1,5 @@
 import ts from "typescript";
+const { factory: f } = ts;
 
 /**
  * Extracts the list of import identifiers from an import clause
@@ -32,3 +33,21 @@ export function getImportIdentifiers(node: ts.ImportDeclaration): string[] {
 
   return importIdentifiers;
 }
+
+/**
+ * Creates an import statement from the given arguments
+ * @param identifiers array of types to import
+ * @param path module path
+ * @returns an ImportDeclaration node that corresponds to `import { ...identifiers } from "path"`
+ */
+export function createImportNode(identifiers: string[], path: string) {
+  const specifiers = identifiers.map((i) =>
+    f.createImportSpecifier(false, undefined, f.createIdentifier(i))
+  );
+
+  return f.createImportDeclaration(
+    undefined,
+    f.createImportClause(false, undefined, f.createNamedImports(specifiers)),
+    f.createStringLiteral(path)
+  );
+}
diff --git a/ts-to-zod.config.js b/ts-to-zod.config.js
index c46d721..ada89fb 100644
--- a/ts-to-zod.config.js
+++ b/ts-to-zod.config.js
@@ -16,5 +16,10 @@ module.exports = [
       },
     },
   },
+  {
+    name: "example/person",
+    input: "example/person.ts",
+    output: "example/person.zod.ts",
+  },
   { name: "config", input: "src/config.ts", output: "src/config.zod.ts" },
 ];


Test patch:
diff --git a/src/core/generate.test.ts b/src/core/generate.test.ts
index 7a3c576..c46043a 100644
--- a/src/core/generate.test.ts
+++ b/src/core/generate.test.ts
@@ -1,3 +1,4 @@
+import { InputOutputMapping } from "../config";
 import { generate } from "./generate";
 
 describe("generate", () => {
@@ -721,5 +722,291 @@ describe("generate", () => {
       `);
       });
     });
+
+    describe("import as union with object literal", () => {
+      const sourceText = `
+      import { Hero } from "module"
+    
+      export type Person = Hero | { name: string }
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+      });
+
+      it("should generate the zod schemas", () => {
+        expect(getZodSchemasFile("./source")).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from "zod";
+
+        const heroSchema = z.any();
+
+        export const personSchema = z.union([heroSchema, z.object({
+                name: z.string()
+            })]);
+        "
+      `);
+      });
+    });
+
+    describe("import as array of union of imported types", () => {
+      const sourceText = `
+      import { Hero } from "module"
+    
+      import { Villain } from "module"
+    
+      export type Person = (Hero | Villain)[]
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+      });
+
+      it("should generate the zod schemas", () => {
+        expect(getZodSchemasFile("./source")).toMatchInlineSnapshot(`
+      "// Generated by ts-to-zod
+      import { z } from "zod";
+
+      const heroSchema = z.any();
+
+      const villainSchema = z.any();
+
+      export const personSchema = z.array(z.union([heroSchema, villainSchema]));
+      "
+      `);
+      });
+    });
+  });
+
+  describe("with input/output mappings to manage imports", () => {
+    describe("one import in one statement", () => {
+      const input = "./hero";
+      const output = "./hero.zod";
+      const inputOutputMappings = [{ input, output }];
+
+      const sourceText = `
+      import { Hero } from "${input}"
+    
+      export interface Person {
+        id: number
+        hero: Hero
+      }
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+        inputOutputMappings,
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from "zod";
+
+        import { heroSchema } from "${output}";
+
+        export const personSchema = z.object({
+            id: z.number(),
+            hero: heroSchema
+        });
+        "
+      `);
+      });
+    });
+
+    describe("multiple imports from one statement", () => {
+      const input = "./hero";
+      const output = "./hero.zod";
+      const inputOutputMappings = [{ input, output }];
+
+      const sourceText = `
+      import { Hero, Villain} from "${input}"
+    
+      export interface Person {
+        id: number
+        hero: Hero
+        villain: Villain
+      }
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+        inputOutputMappings,
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from "zod";
+
+        import { heroSchema, villainSchema } from "${output}";
+
+        export const personSchema = z.object({
+            id: z.number(),
+            hero: heroSchema,
+            villain: villainSchema
+        });
+        "
+      `);
+      });
+    });
+
+    describe("imports from multiple statement", () => {
+      const input = "./hero";
+      const output = "./hero.zod";
+      const input2 = "./villain";
+      const output2 = "./villain.zod";
+      const inputOutputMappings = [
+        { input, output },
+        { input: input2, output: output2 },
+      ];
+
+      const sourceText = `
+      import { Hero } from "${input}"
+      import Villain from "${input2}"
+    
+      export interface Person {
+        id: number
+        hero: Hero
+        villain: Villain
+      }
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+        inputOutputMappings,
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from "zod";
+
+        import { heroSchema } from "${output}";
+        import { villainSchema } from "${output2}";
+
+        export const personSchema = z.object({
+            id: z.number(),
+            hero: heroSchema,
+            villain: villainSchema
+        });
+        "
+      `);
+      });
+    });
+
+    describe("one import in one statement, alternate getSchemaName for mapping", () => {
+      const input = "./hero";
+      const output = "./hero.zod";
+      const inputOutputMappings: InputOutputMapping[] = [
+        { input, output, getSchemaName: (id) => `z${id}` },
+      ];
+
+      const sourceText = `
+      import { Hero } from "${input}"
+    
+      export interface Person {
+        id: number
+        hero: Hero
+      }
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+        inputOutputMappings,
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from "zod";
+
+        import { zHero } from "${output}";
+
+        export const personSchema = z.object({
+            id: z.number(),
+            hero: zHero
+        });
+        "
+      `);
+      });
+    });
+
+    describe("one import in one statement, alternate getSchemaName for generate", () => {
+      const input = "./hero";
+      const output = "./hero.zod";
+      const inputOutputMappings: InputOutputMapping[] = [{ input, output }];
+
+      const sourceText = `
+      import { Hero } from "${input}"
+    
+      export interface Person {
+        id: number
+        hero: Hero
+      }
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+        inputOutputMappings,
+        getSchemaName: (id) => `z${id}`, // should be used only for Person
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from "zod";
+
+        import { heroSchema } from "${output}";
+
+        export const zPerson = z.object({
+            id: z.number(),
+            hero: heroSchema
+        });
+        "
+      `);
+      });
+    });
+
+    describe("mixed imports, 3rd party & zod", () => {
+      const input = "./hero";
+      const output = "./hero.zod";
+      const inputOutputMappings: InputOutputMapping[] = [{ input, output }];
+
+      const sourceText = `
+      import { Hero } from "${input}"
+
+      import Villain from "module"
+    
+      export interface Person {
+        id: number
+        hero: Hero
+        villain: Villain
+      }
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+        inputOutputMappings,
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from "zod";
+        
+        import { heroSchema } from "${output}";
+        
+        const villainSchema = z.any();
+        
+        export const personSchema = z.object({
+            id: z.number(),
+            hero: heroSchema,
+            villain: villainSchema
+        });
+        "
+      `);
+      });
+    });
   });
 });
diff --git a/src/core/validateGeneratedTypes.test.ts b/src/core/validateGeneratedTypes.test.ts
index 6200110..5fe54b0 100644
--- a/src/core/validateGeneratedTypes.test.ts
+++ b/src/core/validateGeneratedTypes.test.ts
@@ -104,6 +104,7 @@ describe("validateGeneratedTypes", () => {
         villain: {
           name: string
           id: any
+          ids: any[]
         };
       };
     `,
@@ -117,6 +118,7 @@ describe("validateGeneratedTypes", () => {
       villain: z.object({
         name: z.string(),
         id: z.any()
+        ids: z.array(z.any())
       })
     });
     `,
diff --git a/src/utils/importHandling.test.ts b/src/utils/importHandling.test.ts
index d65de71..994b9f6 100644
--- a/src/utils/importHandling.test.ts
+++ b/src/utils/importHandling.test.ts
@@ -1,6 +1,6 @@
 import ts from "typescript";
 import { findNode } from "./findNode";
-import { getImportIdentifiers } from "./importHandling";
+import { createImportNode, getImportIdentifiers } from "./importHandling";
 
 describe("getImportIdentifiers", () => {
   it("should return nothing with a StringLiteral import", () => {
@@ -74,6 +74,40 @@ describe("getImportIdentifiers", () => {
   });
 });
 
+describe("createImportNode", () => {
+  function printNode(node: ts.Node) {
+    const printer = ts.createPrinter();
+    return printer.printNode(
+      ts.EmitHint.Unspecified,
+      node,
+      ts.createSourceFile("", "", ts.ScriptTarget.Latest)
+    );
+  }
+
+  it("should create an ImportDeclaration node correctly", () => {
+    const identifiers = ["Test1", "Test2"];
+    const path = "./testPath";
+
+    const expected = 'import { Test1, Test2 } from "./testPath";';
+
+    const result = createImportNode(identifiers, path);
+
+    expect(printNode(result)).toEqual(expected);
+  });
+
+  it("should handle empty identifiers array", () => {
+    const identifiers: string[] = [];
+    const path = "./testPath";
+
+    // Yes, this is valid
+    const expected = 'import {} from "./testPath";';
+
+    const result = createImportNode(identifiers, path);
+
+    expect(printNode(result)).toEqual(expected);
+  });
+});
+
 function getImportNode(sourceText: string): ts.ImportDeclaration {
   const sourceFile = ts.createSourceFile(
     "index.ts",