Merge pull request #17 from DouglasNeuroInformatics/dev

feat: add $$Function helper to replicate Zod v3 behaviour

Author: joshunrau

src/__tests__/zod.test.ts

@@ -3,7 +3,17 @@ import { describe, expect, expectTypeOf, it, test } from 'vitest';
 import { z as z3 } from 'zod/v3';
 import { z as z4 } from 'zod/v4';
 
-import { $BooleanLike, $NumberLike, $Uint8ArrayLike, $UrlLike, isZodType, isZodTypeLike, safeParse } from '../zod.js';
+import {
+  $$Function,
+  $AnyFunction,
+  $BooleanLike,
+  $NumberLike,
+  $Uint8ArrayLike,
+  $UrlLike,
+  isZodType,
+  isZodTypeLike,
+  safeParse
+} from '../zod.js';
 
 import type {
   ZodErrorLike,
@@ -216,6 +226,41 @@ describe('$Uint8ArrayLike', () => {
   });
 });
 
+describe('$AnyFunction', () => {
+  it('should be correctly typed', () => {
+    expectTypeOf<z4.infer<typeof $AnyFunction>>().toEqualTypeOf<(...args: any[]) => any>();
+  });
+  it('should fail to validate a non-function', () => {
+    expect($AnyFunction.safeParse('').success).toBe(false);
+  });
+  it('should validate a function', () => {
+    expect($AnyFunction.safeParse(safeParse).success).toBe(true);
+  });
+});
+
+describe('$$Function', () => {
+  const $Schema = $$Function({ input: [z4.number(), z4.number()], output: z4.number() });
+  it('should be correctly typed', () => {
+    expectTypeOf<z4.infer<typeof $Schema>>().toEqualTypeOf<(arg_0: number, arg_1: number) => number>();
+  });
+  it('should fail to validate a non-function', () => {
+    expect($Schema.safeParse('').success).toBe(false);
+  });
+  it('should return a function that throws when called with the incorrect number of arguments', () => {
+    const fn = $Schema.parse((..._args: any[]) => 0) as (...args: any[]) => number;
+    expect(() => fn(1)).toThrow();
+    expect(() => fn(1, 2, 3)).toThrow();
+  });
+  it('should return a function that throws when it returns an invalid value', () => {
+    const fn = $Schema.parse((..._args: any[]) => 'hello') as (...args: any[]) => any;
+    expect(() => fn(1, 2)).toThrow();
+  });
+  it('should return a function that returns a valid value, when called with valid inputs', () => {
+    const fn = $Schema.parse((a: number, b: number) => a + b);
+    expect(fn(1, 2)).toBe(3);
+  });
+});
+
 describe('safeParse', () => {
   const $Schema = z4.object({ foo: z4.enum(['1', '2']).transform(Number) });
   it('should return an Ok result with the parsed data if successful', () => {

src/zod.ts

@@ -119,6 +119,28 @@ export const $Uint8ArrayLike: z4.ZodType<Uint8Array> = z4
     return arg;
   });
 
+export const $AnyFunction = z4.custom<(...args: any[]) => any>((arg) => typeof arg === 'function', 'must be function');
+
+export const $$Function = <TInput extends [z4.ZodType, ...z4.ZodType[]], TOutput extends z4.ZodType>({
+  input,
+  output
+}: {
+  input: TInput;
+  output: TOutput;
+}): z4.ZodType<(...args: z4.output<z4.ZodTuple<TInput, null>>) => z4.output<TOutput>> => {
+  const $Schema = z4.function({
+    input: z4.tuple(input),
+    output
+  });
+  return z4.custom().transform((arg, ctx) => {
+    if (typeof arg !== 'function') {
+      ctx.addIssue('Must be function');
+      return z4.NEVER;
+    }
+    return $Schema.implement(arg as (...args: any[]) => any);
+  });
+};
+
 export function safeParse<TSchema extends z4.ZodTypeAny>(
   data: unknown,
   $Schema: TSchema