support converting native enum to union (#34)

README.md

@@ -317,7 +317,7 @@ result:
 }
 ```
 
-There are two ways to solve this: provide an identifier to it or resolve all the enums inside `zodToTs()`.
+There are three ways to solve this: provide an identifier to it or resolve all the enums inside `zodToTs()`.
 
 Option 1 - providing an identifier using `withGetType()`:
 
@@ -358,7 +358,7 @@ result:
 Option 2 - resolve enums. This is the same as before, but you just need to pass an option:
 
 ```ts
-const TreeTSType = zodToTs(TreeSchema, undefined, { resolveNativeEnums: true })
+const TreeTSType = zodToTs(TreeSchema, undefined, { nativeEnums: 'resolve' })
 ```
 
 result:
@@ -384,3 +384,22 @@ result:
 Note: These are not the actual values, they are TS representation. The actual values are TS AST nodes.
 
 This option allows you to embed the enums before the schema without actually depending on an external enum type.
+
+Option 3 - convert to union. This is the same as how ZodEnum created by z.enum([...]) is handled, but need to pass an option:
+
+```ts
+const { node } = zodToTs(TreeSchema, undefined, {
+	nativeEnums: 'union',
+})
+```
+
+result:
+
+<!-- dprint-ignore -->
+```ts
+{
+  fruit: 'apple' | 'banana' | 'cantaloupe'
+}
+```
+
+Note: These are not the actual values, they are TS representation. The actual values are TS AST nodes.

src/index.ts

@@ -4,7 +4,8 @@ import {
 	GetType,
 	GetTypeFunction,
 	LiteralType,
-	RequiredZodToTsOptions,
+	ResolvedZodToTsOptions,
+	resolveOptions,
 	ZodToTsOptions,
 	ZodToTsReturn,
 	ZodToTsStore,
@@ -22,7 +23,7 @@ const { factory: f, SyntaxKind } = ts
 const callGetType = (
 	zod: ZodTypeAny & GetType,
 	identifier: string,
-	options: RequiredZodToTsOptions,
+	options: ZodToTsOptions,
 ) => {
 	let type: ReturnType<GetTypeFunction> | undefined
 
@@ -31,11 +32,6 @@ const callGetType = (
 	return type
 }
 
-export const resolveOptions = (raw?: ZodToTsOptions): RequiredZodToTsOptions => {
-	const resolved: RequiredZodToTsOptions = { resolveNativeEnums: true }
-	return { ...resolved, ...raw }
-}
-
 export const zodToTs = (
 	zod: ZodTypeAny,
 	identifier?: string,
@@ -56,7 +52,7 @@ const zodToTsNode = (
 	zod: ZodTypeAny,
 	identifier: string,
 	store: ZodToTsStore,
-	options: RequiredZodToTsOptions,
+	options: ResolvedZodToTsOptions,
 ) => {
 	const typeName = zod._def.typeName
 
@@ -167,7 +163,7 @@ const zodToTsNode = (
 
 		case 'ZodEnum': {
 			// z.enum['a', 'b', 'c'] -> 'a' | 'b' | 'c
-			const types = zod._def.values.map((value: string) => f.createStringLiteral(value))
+			const types = zod._def.values.map((value: string) => f.createLiteralTypeNode(f.createStringLiteral(value)))
 			return f.createUnionTypeNode(types)
 		}
 
@@ -192,12 +188,26 @@ const zodToTsNode = (
 		}
 
 		case 'ZodNativeEnum': {
+			const type = getTypeType
+
+			if (options.nativeEnums === 'union') {
+				// allow overriding with this option
+				if (type) return maybeIdentifierToTypeReference(type)
+
+				const types = Object.values(zod._def.values).map((value) => {
+					if (typeof value === 'number') {
+						return f.createLiteralTypeNode(f.createNumericLiteral(value))
+					}
+					return f.createLiteralTypeNode(f.createStringLiteral(value as string))
+				})
+				return f.createUnionTypeNode(types)
+			}
+
 			// z.nativeEnum(Fruits) -> Fruits
 			// can resolve Fruits into store and user can handle enums
-			let type = getTypeType
 			if (!type) return createUnknownKeywordNode()
 
-			if (options.resolveNativeEnums) {
+			if (options.nativeEnums === 'resolve') {
 				const enumMembers = Object.entries(zod._def.values as Record<string, string | number>).map(([key, value]) => {
 					const literal = typeof value === 'number'
 						? f.createNumericLiteral(value)
@@ -218,13 +228,11 @@ const zodToTsNode = (
 						),
 					)
 				} else {
-					throw new Error('getType on nativeEnum must return an identifier when resolveNativeEnums is set')
+					throw new Error('getType on nativeEnum must return an identifier when nativeEnums is "resolve"')
 				}
 			}
 
-			type = maybeIdentifierToTypeReference(type)
-
-			return type
+			return maybeIdentifierToTypeReference(type)
 		}
 
 		case 'ZodOptional': {

src/types.ts

@@ -3,10 +3,20 @@ import ts from 'typescript'
 export type LiteralType = string | number | boolean
 
 export type ZodToTsOptions = {
+	/** @deprecated use `nativeEnums` instead */
 	resolveNativeEnums?: boolean
+	nativeEnums?: 'identifier' | 'resolve' | 'union'
 }
 
-export type RequiredZodToTsOptions = Required<ZodToTsOptions>
+export const resolveOptions = (raw?: ZodToTsOptions) => {
+	const resolved = {
+		nativeEnums: raw?.resolveNativeEnums ? 'resolve' : 'identifier',
+	} satisfies ZodToTsOptions
+
+	return { ...resolved, ...raw }
+}
+
+export type ResolvedZodToTsOptions = ReturnType<typeof resolveOptions>
 
 export type ZodToTsStore = {
 	nativeEnums: ts.EnumDeclaration[]
@@ -20,7 +30,7 @@ export type ZodToTsReturn = {
 export type GetTypeFunction = (
 	typescript: typeof ts,
 	identifier: string,
-	options: RequiredZodToTsOptions,
+	options: ResolvedZodToTsOptions,
 ) => ts.Identifier | ts.TypeNode
 
 export type GetType = { _def: { getType?: GetTypeFunction } }

test/enum.test.ts

@@ -55,7 +55,7 @@ describe('Fruit enum', () => {
 			(ts) => ts.factory.createIdentifier('Fruit'),
 		)
 
-		const { store } = zodToTs(schema, undefined, { resolveNativeEnums: true })
+		const { store } = zodToTs(schema, undefined, { nativeEnums: 'resolve' })
 
 		expect(printNodeTest(store.nativeEnums[0])).toMatchInlineSnapshot(`
 			"enum Fruit {
@@ -79,7 +79,7 @@ it('handles string literal properties', () => {
 		(ts) => ts.factory.createIdentifier('StringLiteral'),
 	)
 
-	const { store } = zodToTs(schema, undefined, { resolveNativeEnums: true })
+	const { store } = zodToTs(schema, undefined, { nativeEnums: 'resolve' })
 
 	expect(printNodeTest(store.nativeEnums[0])).toMatchInlineSnapshot(`
 		"enum StringLiteral {
@@ -92,3 +92,27 @@ it('handles string literal properties', () => {
 		}"
 	`)
 })
+
+describe('convertNativeEnumToUnion option', () => {
+	it('handles number enum', () => {
+		enum Color {
+			Red,
+			Green,
+			Blue,
+		}
+		const schema = z.nativeEnum(Color)
+		const { node } = zodToTs(schema, undefined, { nativeEnums: 'union' })
+		expect(printNodeTest(node)).toMatchInlineSnapshot(`"\\"Red\\" | \\"Green\\" | \\"Blue\\" | 0 | 1 | 2"`)
+	})
+
+	it('handles string enum', () => {
+		enum Fruit {
+			Apple = 'apple',
+			Banana = 'banana',
+			Cantaloupe = 'cantaloupe',
+		}
+		const schema = z.nativeEnum(Fruit)
+		const { node } = zodToTs(schema, undefined, { nativeEnums: 'union' })
+		expect(printNodeTest(node)).toMatchInlineSnapshot(`"\\"apple\\" | \\"banana\\" | \\"cantaloupe\\""`)
+	})
+})

test/example.ts

@@ -22,7 +22,7 @@ const pickedSchema = example2.partial()
 const nativeEnum = withGetType(z.nativeEnum(Fruits), (ts, _, options) => {
 	const identifier = ts.factory.createIdentifier('Fruits')
 
-	if (options.resolveNativeEnums) return identifier
+	if (options.nativeEnums === 'resolve') return identifier
 
 	return ts.factory.createTypeReferenceNode(
 		identifier,
@@ -103,7 +103,7 @@ type A = z.infer<typeof example>['ee']
 
 type B = z.infer<typeof pickedSchema>
 
-const { node, store } = zodToTs(example, 'Example', { resolveNativeEnums: true })
+const { node, store } = zodToTs(example, 'Example', { nativeEnums: 'resolve' })
 
 console.log(printNode(node))
 

test/example2.ts

@@ -12,7 +12,7 @@ const schema = z.object({
 	key: Enum.describe('Comment for key'),
 })
 
-const { node } = zodToTs(schema, undefined, { resolveNativeEnums: true })
+const { node } = zodToTs(schema, undefined, { nativeEnums: 'resolve' })
 console.log(printNode(node))
 // {
 //     /** Comment for key */