idea: leverage module augmentation to reduce bundle size

[mfulton26]

I like how Zod has great auto-completion, but it currently results in unused code being loaded which increases bundle sizes (e.g. see #2596).

Module augmentation could be a great way to reduce the core Zod size and allow developers to optionally load additional schema methods. Zod can still be a large library. In fact, it can grow as much as wanted with lots of methods. Consumers will still only load what they want. Their bundle sizes will remain small and only grow as they opt-into more methods. This results in faster loading web pages, cloud function (e.g. AWS Lambda), etc.

This could be a breaking change or Zod could continue to export everything together but provide alternative import paths for this module augmentation leveraging approach.

Here is an oversimplified example that only loads desired methods:

import { z } from "./zod";
import "./zod/string";
import "./zod/string/max";

console.log(z.string().max(12).parse("abc")); // => "abc"

// cannot use .number() unless it is loaded by importing ./zod/number"
// console.log(z.number().parse(42));
//               ^^^^^^ error TS2339: Property 'number' does not exist on type 'Zod'.

zod/index.ts

export class Zod {}

export const z = new Zod();

type ZodCheck<T> = (value: T) => void;

export abstract class ZodSchema<T> {
  #checks: ZodCheck<T>[];

  constructor(checks: ZodCheck<T>[] = []) {
    this.#checks = checks;
  }

  protected abstract typeCheck(value: unknown): asserts value is T;

  parse(value: unknown) {
    this.typeCheck(value);
    for (const check of this.#checks) check(value);
    return value;
  }

  refinement(check: ZodCheck<T>): this {
    type Self = this;
    return new (this.constructor as { new (checks: ZodCheck<T>[]): Self })([
      ...this.#checks,
      check,
    ]);
  }
}

zod/number.ts

import type { Zod } from ".";
import { z, ZodSchema } from ".";

export class ZodNumber extends ZodSchema<number> {
  protected typeCheck(value: unknown): asserts value is number {
    if (typeof value === "number") return;
    throw new Error("value is not a number");
  }
}

declare module "." {
  interface Zod {
    number(): ZodNumber;
  }
}

z.number = function (this: Zod) {
  return new ZodNumber();
};

zod/string.ts

import type { Zod } from ".";
import { z, ZodSchema } from ".";

export class ZodString extends ZodSchema<string> {
  protected typeCheck(value: unknown): asserts value is string {
    if (typeof value === "string") return;
    throw new Error("value is not a string");
  }
}

declare module "." {
  interface Zod {
    string(): ZodString;
  }
}

z.string = function (this: Zod) {
  return new ZodString();
};

zod/string/max.ts

import { ZodString } from "../string";

declare module "../string" {
  interface ZodString {
    max(maxLength: number): ZodString;
  }
}

Object.defineProperty(ZodString.prototype, "max", {
  value(this: ZodString, maxLength: number) {
    return this.refinement((value) => {
      if (value.length <= maxLength) return;
      throw new RangeError("max length exceeded");
    });
  },
});