diff --git a/example/heros.ts b/example/heros.ts index 68639a8..d7d65bc 100644 --- a/example/heros.ts +++ b/example/heros.ts @@ -1,6 +1,16 @@ +export enum EnemyPower { + Flight = "flight", + Strength = "strength", + Speed = "speed", +} + +export type SpeedEnemy = { + power: EnemyPower.Speed; +}; + export interface Enemy { name: string; - powers: string[]; + powers: EnemyPower[]; inPrison: boolean; } @@ -13,7 +23,7 @@ export interface Superman { export interface Villain { name: string; - powers: string[]; + powers: EnemyPower[]; friends: Villain[]; canBeTrusted: never; } diff --git a/example/heros.zod.ts b/example/heros.zod.ts index bb46b69..fe7bc1b 100644 --- a/example/heros.zod.ts +++ b/example/heros.zod.ts @@ -1,10 +1,16 @@ // Generated by ts-to-zod import { z } from "zod"; -import { Villain } from "./heros"; +import { EnemyPower, Villain } from "./heros"; + +export const enemyPowerSchema = z.nativeEnum(EnemyPower); + +export const speedEnemySchema = z.object({ + power: z.literal(EnemyPower.Speed), +}); export const enemySchema = z.object({ name: z.string(), - powers: z.array(z.string()), + powers: z.array(enemyPowerSchema), inPrison: z.boolean(), }); @@ -22,7 +28,7 @@ export const supermanSchema = z.object({ export const villainSchema: z.ZodSchema = z.lazy(() => z.object({ name: z.string(), - powers: z.array(z.string()), + powers: z.array(enemyPowerSchema), friends: z.array(villainSchema), canBeTrusted: z.never(), }) diff --git a/src/core/generate.test.ts b/src/core/generate.test.ts index 2ea2188..8b2e386 100644 --- a/src/core/generate.test.ts +++ b/src/core/generate.test.ts @@ -72,6 +72,66 @@ describe("generate", () => { " `); }); + it("should not have any errors", () => { + expect(errors.length).toBe(0); + }); + }); + + describe("with enums", () => { + const sourceText = ` + export enum Superhero { + Superman = "superman" + ClarkKent = "clark-kent" + }; + + export type FavoriteSuperhero = { + superhero: Superhero.Superman + }; + `; + + const { getZodSchemasFile, getIntegrationTestFile, errors } = generate({ + sourceText, + }); + + it("should generate the zod schemas", () => { + expect(getZodSchemasFile("./superhero")).toMatchInlineSnapshot(` + "// Generated by ts-to-zod + import { z } from \\"zod\\"; + import { Superhero } from \\"./superhero\\"; + + export const superheroSchema = z.nativeEnum(Superhero); + + export const favoriteSuperheroSchema = z.object({ + superhero: z.literal(Superhero.Superman) + }); + " + `); + }); + + it("should generate the integration tests", () => { + expect(getIntegrationTestFile("./superhero", "superhero.zod")) + .toMatchInlineSnapshot(` + "// Generated by ts-to-zod + import { z } from \\"zod\\"; + + import * as spec from \\"./superhero\\"; + import * as generated from \\"superhero.zod\\"; + + // eslint-disable-next-line @typescript-eslint/no-unused-vars + function expectType(_: T) { + /* noop */ + } + + export type superheroSchemaInferredType = z.infer; + + export type favoriteSuperheroSchemaInferredType = z.infer; + expectType({} as superheroSchemaInferredType) + expectType({} as spec.Superhero) + expectType({} as favoriteSuperheroSchemaInferredType) + expectType({} as spec.FavoriteSuperhero) + " + `); + }); it("should not have any errors", () => { expect(errors.length).toBe(0); diff --git a/src/core/generate.ts b/src/core/generate.ts index 1a6bd00..533c7b5 100644 --- a/src/core/generate.ts +++ b/src/core/generate.ts @@ -53,10 +53,16 @@ export function generate({ ); // Extract the nodes (interface declarations & type aliases) - const nodes: Array = []; + const nodes: Array< + ts.InterfaceDeclaration | ts.TypeAliasDeclaration | ts.EnumDeclaration + > = []; const visitor = (node: ts.Node) => { - if (ts.isInterfaceDeclaration(node) || ts.isTypeAliasDeclaration(node)) { + if ( + ts.isInterfaceDeclaration(node) || + ts.isTypeAliasDeclaration(node) || + ts.isEnumDeclaration(node) + ) { if (nameFilter(node.name.text)) { nodes.push(node); } @@ -91,23 +97,28 @@ export function generate({ while (statements.size !== zodSchemas.length && n < maxRun) { zodSchemas .filter(({ varName }) => !statements.has(varName)) - .forEach(({ varName, dependencies, statement, typeName }) => { - const isCircular = dependencies.includes(varName); - const missingDependencies = dependencies - .filter((dep) => dep !== varName) - .filter((dep) => !statements.has(dep)); - if (missingDependencies.length === 0) { - if (isCircular) { - typeImports.add(typeName); - statements.set(varName, { - value: transformRecursiveSchema("z", statement, typeName), - typeName, - }); - } else { - statements.set(varName, { value: statement, typeName }); + .forEach( + ({ varName, dependencies, statement, typeName, requiresImport }) => { + const isCircular = dependencies.includes(varName); + const missingDependencies = dependencies + .filter((dep) => dep !== varName) + .filter((dep) => !statements.has(dep)); + if (missingDependencies.length === 0) { + if (isCircular) { + typeImports.add(typeName); + statements.set(varName, { + value: transformRecursiveSchema("z", statement, typeName), + typeName, + }); + } else { + if (requiresImport) { + typeImports.add(typeName); + } + statements.set(varName, { value: statement, typeName }); + } } } - }); + ); n++; // Just a safety net to avoid infinity loops } diff --git a/src/core/generateZodSchema.test.ts b/src/core/generateZodSchema.test.ts index 9ea8d28..25064b7 100644 --- a/src/core/generateZodSchema.test.ts +++ b/src/core/generateZodSchema.test.ts @@ -95,6 +95,29 @@ describe("generateZodSchema", () => { ); }); + it("should generate a literal schema (enum)", () => { + const source = ` + export type BestSuperhero = { + superhero: Superhero.Superman + }; + `; + expect(generate(source)).toMatchInlineSnapshot(` + "export const bestSuperheroSchema = z.object({ + superhero: z.literal(Superhero.Superman) + });" + `); + }); + + it("should generate a nativeEnum schema", () => { + const source = `export enum Superhero = { + Superman = "superman", + ClarkKent = "clark_kent", + };`; + expect(generate(source)).toMatchInlineSnapshot( + `"export const superheroSchema = z.nativeEnum(Superhero);"` + ); + }); + it("should generate a never", () => { const source = `export type CanBeatZod = never;`; expect(generate(source)).toMatchInlineSnapshot( @@ -620,8 +643,15 @@ function generate(sourceText: string, z?: string) { ); const declaration = findNode( sourceFile, - (node): node is ts.InterfaceDeclaration | ts.TypeAliasDeclaration => - ts.isInterfaceDeclaration(node) || ts.isTypeAliasDeclaration(node) + ( + node + ): node is + | ts.InterfaceDeclaration + | ts.TypeAliasDeclaration + | ts.EnumDeclaration => + ts.isInterfaceDeclaration(node) || + ts.isTypeAliasDeclaration(node) || + ts.isEnumDeclaration(node) ); if (!declaration) { throw new Error("No `type` or `interface` found!"); diff --git a/src/core/generateZodSchema.ts b/src/core/generateZodSchema.ts index b0311f3..cdd58ed 100644 --- a/src/core/generateZodSchema.ts +++ b/src/core/generateZodSchema.ts @@ -19,7 +19,7 @@ export interface GenerateZodSchemaProps { /** * Interface or type node */ - node: ts.InterfaceDeclaration | ts.TypeAliasDeclaration; + node: ts.InterfaceDeclaration | ts.TypeAliasDeclaration | ts.EnumDeclaration; /** * Zod import value. @@ -57,6 +57,7 @@ export function generateZodSchemaVariableStatement({ }: GenerateZodSchemaProps) { let schema: ts.CallExpression | ts.Identifier | undefined; const dependencies: string[] = []; + let requiresImport = false; if (ts.isInterfaceDeclaration(node)) { let baseSchema: string | undefined; @@ -100,6 +101,11 @@ export function generateZodSchemaVariableStatement({ }); } + if (ts.isEnumDeclaration(node)) { + schema = buildZodSchema(zodImportValue, "nativeEnum", [node.name]); + requiresImport = true; + } + return { dependencies: uniq(dependencies), statement: f.createVariableStatement( @@ -116,6 +122,7 @@ export function generateZodSchemaVariableStatement({ ts.NodeFlags.Const ) ), + requiresImport, }; } @@ -447,6 +454,25 @@ function buildZodPrimitive({ } } + // Deal with enums used as literals + if ( + ts.isTypeReferenceNode(typeNode) && + ts.isQualifiedName(typeNode.typeName) && + ts.isIdentifier(typeNode.typeName.left) + ) { + return buildZodSchema( + z, + "literal", + [ + f.createPropertyAccessExpression( + typeNode.typeName.left, + typeNode.typeName.right + ), + ], + zodProperties + ); + } + if (ts.isArrayTypeNode(typeNode)) { return buildZodSchema( z,