Fix inferred function types

colinhacks · Jul 18, 2022 · a59c384 · a59c384
1 parent fb6b0a7
commit a59c384
Show file tree

Hide file tree

Showing 8 changed files with 84 additions and 28 deletions.
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -111,7 +111,7 @@ jobs:
         with:
           tag_name: v${{ steps.publish.outputs.version }}
           release_name: v${{ steps.publish.outputs.version }}
-          commitish: ref/heads/main
+          commitish: refs/heads/main
           body: ${{steps.github_release.outputs.changelog}}
           draft: false
           prerelease: false

diff --git a/deno/lib/__tests__/function.test.ts b/deno/lib/__tests__/function.test.ts
@@ -26,26 +26,21 @@ test("parsed function fail 2", () => {
 
 test("function inference 1", () => {
   type func1 = z.TypeOf<typeof func1>;
-  const t1: util.AssertEqual<func1, (k: string) => number> = true;
-  [t1];
+  util.assertEqual<func1, (k: string) => number>(true);
 });
 
 test("args method", () => {
   const t1 = z.function();
   type t1 = z.infer<typeof t1>;
-  const f1: util.AssertEqual<t1, () => void> = true;
+  util.assertEqual<t1, () => void>(true);
 
   const t2 = t1.args(z.string());
   type t2 = z.infer<typeof t2>;
-  const f2: util.AssertEqual<t2, (arg: string) => void> = true;
+  util.assertEqual<t2, (arg: string) => void>(true);
 
   const t3 = t2.returns(z.boolean());
   type t3 = z.infer<typeof t3>;
-  const f3: util.AssertEqual<t3, (arg: string) => boolean> = true;
-
-  f1;
-  f2;
-  f3;
+  util.assertEqual<t3, (arg: string) => boolean>(true);
 });
 
 const args2 = z.tuple([
@@ -61,15 +56,14 @@ const func2 = z.function(args2, returns2);
 
 test("function inference 2", () => {
   type func2 = z.TypeOf<typeof func2>;
-  const t2: util.AssertEqual<
+  util.assertEqual<
     func2,
     (arg: {
       f1: number;
       f2: string | null;
       f3?: (boolean | undefined)[] | undefined;
     }) => string | number
-  > = true;
-  [t2];
+  >(true);
 });
 
 test("valid function run", () => {
@@ -212,3 +206,32 @@ test("params and returnType getters", () => {
   func.parameters().items[0].parse("asdf");
   func.returnType().parse("asdf");
 });
+
+test("inference with transforms", () => {
+  const funcSchema = z
+    .function()
+    .args(z.string().transform((val) => val.length))
+    .returns(z.object({ val: z.number() }));
+  const myFunc = funcSchema.implement((val) => {
+    return { val, extra: "stuff" };
+  });
+  myFunc("asdf");
+
+  util.assertEqual<
+    typeof myFunc,
+    (arg: string) => { val: number; extra: string }
+  >(true);
+});
+
+test("fallback to OuterTypeOfFunction", () => {
+  const funcSchema = z
+    .function()
+    .args(z.string().transform((val) => val.length))
+    .returns(z.object({ arg: z.number() }).transform((val) => val.arg));
+
+  const myFunc = funcSchema.implement((val) => {
+    return { arg: val, arg2: false };
+  });
+
+  util.assertEqual<typeof myFunc, (arg: string) => number>(true);
+});
diff --git a/deno/lib/helpers/util.ts b/deno/lib/helpers/util.ts
@@ -4,6 +4,7 @@ export namespace util {
       ? true
       : false
     : false;
+  export function assertEqual<A, B>(_cond: AssertEqual<A, B>) {}
 
   export function assertNever(_x: never): never {
     throw new Error();

diff --git a/deno/lib/types.ts b/deno/lib/types.ts
@@ -2915,7 +2915,11 @@ export class ZodFunction<
     });
   }
 
-  implement<F extends InnerTypeOfFunction<Args, Returns>>(func: F): F {
+  implement<F extends InnerTypeOfFunction<Args, Returns>>(
+    func: F
+  ): ReturnType<F> extends Returns["_output"]
+    ? (...args: Args["_input"]) => ReturnType<F>
+    : OuterTypeOfFunction<Args, Returns> {
     const validatedFunc = this.parse(func);
     return validatedFunc as any;
   }

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "zod",
-  "version": "3.17.6",
+  "version": "3.17.7",
   "description": "TypeScript-first schema declaration and validation library with static type inference",
   "main": "./lib/index.js",
   "types": "./index.d.ts",

diff --git a/src/__tests__/function.test.ts b/src/__tests__/function.test.ts
@@ -25,26 +25,21 @@ test("parsed function fail 2", () => {
 
 test("function inference 1", () => {
   type func1 = z.TypeOf<typeof func1>;
-  const t1: util.AssertEqual<func1, (k: string) => number> = true;
-  [t1];
+  util.assertEqual<func1, (k: string) => number>(true);
 });
 
 test("args method", () => {
   const t1 = z.function();
   type t1 = z.infer<typeof t1>;
-  const f1: util.AssertEqual<t1, () => void> = true;
+  util.assertEqual<t1, () => void>(true);
 
   const t2 = t1.args(z.string());
   type t2 = z.infer<typeof t2>;
-  const f2: util.AssertEqual<t2, (arg: string) => void> = true;
+  util.assertEqual<t2, (arg: string) => void>(true);
 
   const t3 = t2.returns(z.boolean());
   type t3 = z.infer<typeof t3>;
-  const f3: util.AssertEqual<t3, (arg: string) => boolean> = true;
-
-  f1;
-  f2;
-  f3;
+  util.assertEqual<t3, (arg: string) => boolean>(true);
 });
 
 const args2 = z.tuple([
@@ -60,15 +55,14 @@ const func2 = z.function(args2, returns2);
 
 test("function inference 2", () => {
   type func2 = z.TypeOf<typeof func2>;
-  const t2: util.AssertEqual<
+  util.assertEqual<
     func2,
     (arg: {
       f1: number;
       f2: string | null;
       f3?: (boolean | undefined)[] | undefined;
     }) => string | number
-  > = true;
-  [t2];
+  >(true);
 });
 
 test("valid function run", () => {
@@ -211,3 +205,32 @@ test("params and returnType getters", () => {
   func.parameters().items[0].parse("asdf");
   func.returnType().parse("asdf");
 });
+
+test("inference with transforms", () => {
+  const funcSchema = z
+    .function()
+    .args(z.string().transform((val) => val.length))
+    .returns(z.object({ val: z.number() }));
+  const myFunc = funcSchema.implement((val) => {
+    return { val, extra: "stuff" };
+  });
+  myFunc("asdf");
+
+  util.assertEqual<
+    typeof myFunc,
+    (arg: string) => { val: number; extra: string }
+  >(true);
+});
+
+test("fallback to OuterTypeOfFunction", () => {
+  const funcSchema = z
+    .function()
+    .args(z.string().transform((val) => val.length))
+    .returns(z.object({ arg: z.number() }).transform((val) => val.arg));
+
+  const myFunc = funcSchema.implement((val) => {
+    return { arg: val, arg2: false };
+  });
+
+  util.assertEqual<typeof myFunc, (arg: string) => number>(true);
+});
diff --git a/src/helpers/util.ts b/src/helpers/util.ts
@@ -4,6 +4,7 @@ export namespace util {
       ? true
       : false
     : false;
+  export function assertEqual<A, B>(_cond: AssertEqual<A, B>) {}
 
   export function assertNever(_x: never): never {
     throw new Error();

diff --git a/src/types.ts b/src/types.ts
@@ -2915,7 +2915,11 @@ export class ZodFunction<
     });
   }
 
-  implement<F extends InnerTypeOfFunction<Args, Returns>>(func: F): F {
+  implement<F extends InnerTypeOfFunction<Args, Returns>>(
+    func: F
+  ): ReturnType<F> extends Returns["_output"]
+    ? (...args: Args["_input"]) => ReturnType<F>
+    : OuterTypeOfFunction<Args, Returns> {
     const validatedFunc = this.parse(func);
     return validatedFunc as any;
   }