From 813872d43d78a9509349733e1240b9c702deb786 Mon Sep 17 00:00:00 2001
From: Noah Davey <Noah.davey@gmail.com>
Date: Wed, 27 Apr 2022 09:43:08 +1000
Subject: [PATCH] Better Optional Handling (#11)

* Add question token for optional properties

* Update tests to reflect correct handling of optionals

* update snapshots

Co-authored-by: Sachin Raja <sachinraja2349@gmail.com>
---
 src/index.ts           |  8 ++++++--
 test/modifiers.test.ts | 39 +++++++++++++++++++++++++++++++++++++--
 test/primitive.test.ts |  8 ++++----
 3 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index d3d7055..816da0f 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -85,12 +85,16 @@ const zodToTsNode = (
       const properties = Object.entries(zod._def.shape())
 
       const members: ts.TypeElement[] = properties.map(([key, value]) => {
-        const type = zodToTsNode(value as ZodTypeAny, ...otherArgs)
+        const nextZodNode = value as ZodTypeAny
+        const type = zodToTsNode(nextZodNode, ...otherArgs)
+
+        const { typeName: nextZodNodeTypeName } = nextZodNode._def
+        const isOptional = nextZodNodeTypeName === 'ZodOptional' || nextZodNode.isOptional()
 
         return f.createPropertySignature(
           undefined,
           f.createIdentifier(key),
-          undefined,
+          isOptional ? f.createToken(ts.SyntaxKind.QuestionToken) : undefined,
           type,
         )
       })
diff --git a/test/modifiers.test.ts b/test/modifiers.test.ts
index 94b7f77..47027df 100644
--- a/test/modifiers.test.ts
+++ b/test/modifiers.test.ts
@@ -5,12 +5,47 @@ import { printNodeTest } from './utils'
 
 const OptionalStringSchema = z.string().optional()
 
-describe('z.optional()', () => {
-  const { node } = zodToTs(OptionalStringSchema)
+const ObjectWithOptionals = z.object({
+  optional: OptionalStringSchema,
+  required: z.string(),
+  transform: z.number().optional().transform((arg) => arg),
+  or: z.number().optional().or(z.string()),
+  tuple: z.tuple([
+    z.string().optional(),
+    z.number(),
+    z.object({
+      optional: z.string().optional(),
+      required: z.string(),
+    }),
+  ]).optional(),
+})
 
+describe('z.optional()', () => {
   it('outputs correct typescript', () => {
+    const { node } = zodToTs(OptionalStringSchema)
     expect(printNodeTest(node)).toMatchInlineSnapshot('"string | undefined"')
   })
+
+  it('for optionals should output ?: property as well as undefined union', () => {
+    const { node } = zodToTs(ObjectWithOptionals)
+
+    expect(printNodeTest(node)).toMatchInlineSnapshot(`
+      "{
+          optional?: string | undefined;
+          required: string;
+          transform?: number | undefined;
+          or?: (number | undefined) | string;
+          tuple?: [
+              string | undefined,
+              number,
+              {
+                  optional?: string | undefined;
+                  required: string;
+              }
+          ] | undefined;
+      }"
+    `)
+  })
 })
 
 const NullableUsernameSchema = z.object({
diff --git a/test/primitive.test.ts b/test/primitive.test.ts
index 1c51409..23e50d2 100644
--- a/test/primitive.test.ts
+++ b/test/primitive.test.ts
@@ -26,11 +26,11 @@ describe('PrimitiveSchema', () => {
           age: number;
           isAdmin: boolean;
           createdAt: Date;
-          undef: undefined;
+          undef?: undefined;
           nu: null;
-          vo: void | undefined;
-          an: any;
-          unknow: unknown;
+          vo?: void | undefined;
+          an?: any;
+          unknow?: unknown;
           nev: never;
       }"
     `)