From 1dd44a0d6f8073f7c417e09ec96580b9ae9bda23 Mon Sep 17 00:00:00 2001
From: Colin McDonnell <colinmcd94@gmail.com>
Date: Tue, 11 Feb 2025 13:43:50 -0800
Subject: [PATCH] Support async z.custom

---
 deno/lib/__tests__/custom.test.ts | 13 +++++++++++
 deno/lib/types.ts                 | 38 ++++++++++++++++++++++---------
 src/__tests__/custom.test.ts      | 13 +++++++++++
 src/types.ts                      | 38 ++++++++++++++++++++++---------
 4 files changed, 80 insertions(+), 22 deletions(-)

diff --git a/deno/lib/__tests__/custom.test.ts b/deno/lib/__tests__/custom.test.ts
index a19233ecf..26a3e50ea 100644
--- a/deno/lib/__tests__/custom.test.ts
+++ b/deno/lib/__tests__/custom.test.ts
@@ -17,3 +17,16 @@ test("string params", () => {
   // @ts-ignore
   expect(JSON.stringify(result.error).includes("customerr")).toEqual(true);
 });
+
+test("async validations", async () => {
+  const example1 = z.custom<number>(async (x) => {
+    return typeof x === "number";
+  });
+  const r1 = await example1.safeParseAsync(1234);
+  expect(r1.success).toEqual(true);
+  expect(r1.data).toEqual(1234);
+
+  const r2 = await example1.safeParseAsync("asdf");
+  expect(r2.success).toEqual(false);
+  expect(r2.error!.issues.length).toEqual(1);
+});
diff --git a/deno/lib/types.ts b/deno/lib/types.ts
index cd09d4b15..acc0c5f01 100644
--- a/deno/lib/types.ts
+++ b/deno/lib/types.ts
@@ -5235,10 +5235,21 @@ export class ZodReadonly<T extends ZodTypeAny> extends ZodType<
 //////////                    //////////
 ////////////////////////////////////////
 ////////////////////////////////////////
+function cleanParams(params: unknown, data: unknown) {
+  const p =
+    typeof params === "function"
+      ? params(data)
+      : typeof params === "string"
+      ? { message: params }
+      : params;
+
+  const p2 = typeof p === "string" ? { message: p } : p;
+  return p2;
+}
 type CustomParams = CustomErrorParams & { fatal?: boolean };
 export function custom<T>(
   check?: (data: any) => any,
-  params: string | CustomParams | ((input: any) => CustomParams) = {},
+  _params: string | CustomParams | ((input: any) => CustomParams) = {},
   /**
    * @deprecated
    *
@@ -5253,17 +5264,22 @@ export function custom<T>(
 ): ZodType<T, ZodTypeDef, T> {
   if (check)
     return ZodAny.create().superRefine((data, ctx) => {
-      if (!check(data)) {
-        const p =
-          typeof params === "function"
-            ? params(data)
-            : typeof params === "string"
-            ? { message: params }
-            : params;
-        const _fatal = p.fatal ?? fatal ?? true;
-        const p2 = typeof p === "string" ? { message: p } : p;
-        ctx.addIssue({ code: "custom", ...p2, fatal: _fatal });
+      const r = check(data);
+      if (r instanceof Promise) {
+        return r.then((r) => {
+          if (!r) {
+            const params = cleanParams(_params, data);
+            const _fatal = params.fatal ?? fatal ?? true;
+            ctx.addIssue({ code: "custom", ...params, fatal: _fatal });
+          }
+        });
+      }
+      if (!r) {
+        const params = cleanParams(_params, data);
+        const _fatal = params.fatal ?? fatal ?? true;
+        ctx.addIssue({ code: "custom", ...params, fatal: _fatal });
       }
+      return;
     });
   return ZodAny.create();
 }
diff --git a/src/__tests__/custom.test.ts b/src/__tests__/custom.test.ts
index 0fb9ca46c..07c698b54 100644
--- a/src/__tests__/custom.test.ts
+++ b/src/__tests__/custom.test.ts
@@ -16,3 +16,16 @@ test("string params", () => {
   // @ts-ignore
   expect(JSON.stringify(result.error).includes("customerr")).toEqual(true);
 });
+
+test("async validations", async () => {
+  const example1 = z.custom<number>(async (x) => {
+    return typeof x === "number";
+  });
+  const r1 = await example1.safeParseAsync(1234);
+  expect(r1.success).toEqual(true);
+  expect(r1.data).toEqual(1234);
+
+  const r2 = await example1.safeParseAsync("asdf");
+  expect(r2.success).toEqual(false);
+  expect(r2.error!.issues.length).toEqual(1);
+});
diff --git a/src/types.ts b/src/types.ts
index 98281ff2f..334473c1b 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -5235,10 +5235,21 @@ export class ZodReadonly<T extends ZodTypeAny> extends ZodType<
 //////////                    //////////
 ////////////////////////////////////////
 ////////////////////////////////////////
+function cleanParams(params: unknown, data: unknown) {
+  const p =
+    typeof params === "function"
+      ? params(data)
+      : typeof params === "string"
+      ? { message: params }
+      : params;
+
+  const p2 = typeof p === "string" ? { message: p } : p;
+  return p2;
+}
 type CustomParams = CustomErrorParams & { fatal?: boolean };
 export function custom<T>(
   check?: (data: any) => any,
-  params: string | CustomParams | ((input: any) => CustomParams) = {},
+  _params: string | CustomParams | ((input: any) => CustomParams) = {},
   /**
    * @deprecated
    *
@@ -5253,17 +5264,22 @@ export function custom<T>(
 ): ZodType<T, ZodTypeDef, T> {
   if (check)
     return ZodAny.create().superRefine((data, ctx) => {
-      if (!check(data)) {
-        const p =
-          typeof params === "function"
-            ? params(data)
-            : typeof params === "string"
-            ? { message: params }
-            : params;
-        const _fatal = p.fatal ?? fatal ?? true;
-        const p2 = typeof p === "string" ? { message: p } : p;
-        ctx.addIssue({ code: "custom", ...p2, fatal: _fatal });
+      const r = check(data);
+      if (r instanceof Promise) {
+        return r.then((r) => {
+          if (!r) {
+            const params = cleanParams(_params, data);
+            const _fatal = params.fatal ?? fatal ?? true;
+            ctx.addIssue({ code: "custom", ...params, fatal: _fatal });
+          }
+        });
+      }
+      if (!r) {
+        const params = cleanParams(_params, data);
+        const _fatal = params.fatal ?? fatal ?? true;
+        ctx.addIssue({ code: "custom", ...params, fatal: _fatal });
       }
+      return;
     });
   return ZodAny.create();
 }