From 0e00f1ea064a3ee01e66ca92260e9adf98407496 Mon Sep 17 00:00:00 2001
From: Fabien BERNARD <fabien0102@hotmail.com>
Date: Mon, 22 Nov 2021 10:27:19 +0100
Subject: [PATCH] feat: Improve nullable (#57)

* Add isNotNull util

* Use `nullable()`
---
 src/core/generateZodSchema.test.ts | 25 ++++++++++++++++++-
 src/core/generateZodSchema.ts      | 39 ++++++++++++++++++++++++------
 src/core/jsDocTags.ts              |  9 ++++++-
 src/utils/isNotNull.ts             | 14 +++++++++++
 4 files changed, 78 insertions(+), 9 deletions(-)
 create mode 100644 src/utils/isNotNull.ts

diff --git a/src/core/generateZodSchema.test.ts b/src/core/generateZodSchema.test.ts
index 414c5cb..3e9cf11 100644
--- a/src/core/generateZodSchema.test.ts
+++ b/src/core/generateZodSchema.test.ts
@@ -357,7 +357,7 @@ describe("generateZodSchema", () => {
           enemies: z.record(enemySchema),
           age: z.number(),
           underKryptonite: z.boolean().optional(),
-          needGlasses: z.union([z.literal(true), z.null()])
+          needGlasses: z.literal(true).nullable()
       });"
     `);
   });
@@ -643,6 +643,29 @@ describe("generateZodSchema", () => {
     `);
   });
 
+  it("should deal with nullable", () => {
+    const source = `export interface A {
+      /** @minimum 0 */
+      a: number | null;
+      /** @minLength 1 */
+      b: string | null;
+      /** @pattern ^c$ */
+      c: string | null;
+    }
+    `;
+
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "export const aSchema = z.object({
+          /** @minimum 0 */
+          a: z.number().min(0).nullable(),
+          /** @minLength 1 */
+          b: z.string().min(1).nullable(),
+          /** @pattern ^c$ */
+          c: z.string().regex(/^c$/).nullable()
+      });"
+    `);
+  });
+
   it("should deal with @default with all types", () => {
     const source = `export interface WithDefaults {
      /**
diff --git a/src/core/generateZodSchema.ts b/src/core/generateZodSchema.ts
index 7210a32..0a39819 100644
--- a/src/core/generateZodSchema.ts
+++ b/src/core/generateZodSchema.ts
@@ -8,6 +8,7 @@ import {
 } from "./jsDocTags";
 import uniq from "lodash/uniq";
 import { findNode } from "../utils/findNode";
+import { isNotNull } from "../utils/isNotNull";
 
 const { factory: f } = ts;
 
@@ -180,6 +181,7 @@ function buildZodPrimitive({
   z,
   typeNode,
   isOptional,
+  isNullable,
   isPartial,
   isRequired,
   jsDocTags,
@@ -190,6 +192,7 @@ function buildZodPrimitive({
   z: string;
   typeNode: ts.TypeNode;
   isOptional: boolean;
+  isNullable?: boolean;
   isPartial?: boolean;
   isRequired?: boolean;
   jsDocTags: JSDocTags;
@@ -201,7 +204,8 @@ function buildZodPrimitive({
     jsDocTags,
     isOptional,
     Boolean(isPartial),
-    Boolean(isRequired)
+    Boolean(isRequired),
+    Boolean(isNullable)
   );
 
   if (ts.isParenthesizedTypeNode(typeNode)) {
@@ -395,22 +399,43 @@ function buildZodPrimitive({
   }
 
   if (ts.isUnionTypeNode(typeNode)) {
-    const values = typeNode.types.map((i) =>
+    const hasNull = Boolean(
+      typeNode.types.find(
+        (i) =>
+          ts.isLiteralTypeNode(i) &&
+          i.literal.kind === ts.SyntaxKind.NullKeyword
+      )
+    );
+
+    const nodes = typeNode.types.filter(isNotNull);
+
+    // type A = | 'b' is a valid typescript definition
+    // Zod does not allow `z.union(['b']), so we have to return just the value
+    if (nodes.length === 1) {
+      return buildZodPrimitive({
+        z,
+        typeNode: nodes[0],
+        isOptional: false,
+        isNullable: hasNull,
+        jsDocTags,
+        sourceFile,
+        dependencies,
+        getDependencyName,
+      });
+    }
+
+    const values = nodes.map((i) =>
       buildZodPrimitive({
         z,
         typeNode: i,
         isOptional: false,
+        isNullable: false,
         jsDocTags: {},
         sourceFile,
         dependencies,
         getDependencyName,
       })
     );
-    // type A = | 'b' is a valid typescript definintion
-    // Zod does not allow `z.union(['b']), so we have to return just the value
-    if (values.length === 1) {
-      return values[0];
-    }
     return buildZodSchema(
       z,
       "union",
diff --git a/src/core/jsDocTags.ts b/src/core/jsDocTags.ts
index aa223ac..4283f61 100644
--- a/src/core/jsDocTags.ts
+++ b/src/core/jsDocTags.ts
@@ -133,12 +133,14 @@ export type ZodProperty = {
  * @param isOptional
  * @param isPartial
  * @param isRequired
+ * @param isNullable
  */
 export function jsDocTagToZodProperties(
   jsDocTags: JSDocTags,
   isOptional: boolean,
   isPartial: boolean,
-  isRequired: boolean
+  isRequired: boolean,
+  isNullable: boolean
 ) {
   const zodProperties: ZodProperty[] = [];
   if (jsDocTags.minimum !== undefined) {
@@ -181,6 +183,11 @@ export function jsDocTagToZodProperties(
       identifier: "optional",
     });
   }
+  if (isNullable) {
+    zodProperties.push({
+      identifier: "nullable",
+    });
+  }
   if (isPartial) {
     zodProperties.push({
       identifier: "partial",
diff --git a/src/utils/isNotNull.ts b/src/utils/isNotNull.ts
new file mode 100644
index 0000000..eb3a7a8
--- /dev/null
+++ b/src/utils/isNotNull.ts
@@ -0,0 +1,14 @@
+import ts from "typescript";
+
+/**
+ * Helper to filter out any `null` node
+ *
+ * @param node
+ * @returns
+ */
+export function isNotNull(node: ts.TypeNode) {
+  return (
+    !ts.isLiteralTypeNode(node) ||
+    node.literal.kind !== ts.SyntaxKind.NullKeyword
+  );
+}