Zod deep strict and `deepPartial` utility

zod-deep-strict-partial.ts

/*
   Copyright 2024, Jaen - https://github.com/jaens

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

import { z, type ZodDiscriminatedUnionOption } from "zod";

const RESOLVING = Symbol("mapOnSchema/resolving");

export function mapOnSchema<T extends z.ZodTypeAny, TResult extends z.ZodTypeAny>(
    schema: T,
    fn: (schema: z.ZodTypeAny) => TResult,
): TResult;

/** 
 * Applies {@link fn} to each element of the schema recursively, replacing every schema with its return value.
 * The rewriting is applied bottom-up (ie. {@link fn} will get called on "children" first).
 */
export function mapOnSchema(schema: z.ZodTypeAny, fn: (schema: z.ZodTypeAny) => z.ZodTypeAny): z.ZodTypeAny {
    // Cache results to support recursive schemas
    const results = new Map<z.ZodTypeAny, z.ZodTypeAny | typeof RESOLVING>();

    function mapElement(s: z.ZodTypeAny) {
        const value = results.get(s);
        if (value === RESOLVING) {
            throw new Error("Recursive schema access detected");
        } else if (value !== undefined) {
            return value;
        }

        results.set(s, RESOLVING);
        const result = mapOnSchema(s, fn);
        results.set(s, result);
        return result;
    }

    function mapInner() {
        if (schema instanceof z.ZodObject) {
            const newShape: Record<string, z.ZodTypeAny> = {};
            for (const [key, value] of Object.entries(schema.shape)) {
                newShape[key] = mapElement(value as z.ZodTypeAny);
            }

            return new z.ZodObject({
                ...schema._def,
                shape: () => newShape,
            });
        } else if (schema instanceof z.ZodArray) {
            return new z.ZodArray({
                ...schema._def,
                type: mapElement(schema._def.type),
            });
        } else if (schema instanceof z.ZodMap) {
            return new z.ZodMap({
                ...schema._def,
                keyType: mapElement(schema._def.keyType),
                valueType: mapElement(schema._def.valueType),
            });
        } else if (schema instanceof z.ZodSet) {
            return new z.ZodSet({
                ...schema._def,
                valueType: mapElement(schema._def.valueType),
            });
        } else if (schema instanceof z.ZodOptional) {
            return new z.ZodOptional({
                ...schema._def,
                innerType: mapElement(schema._def.innerType),
            });
        } else if (schema instanceof z.ZodNullable) {
            return new z.ZodNullable({
                ...schema._def,
                innerType: mapElement(schema._def.innerType),
            });
        } else if (schema instanceof z.ZodDefault) {
            return new z.ZodDefault({
                ...schema._def,
                innerType: mapElement(schema._def.innerType),
            });
        } else if (schema instanceof z.ZodReadonly) {
            return new z.ZodReadonly({
                ...schema._def,
                innerType: mapElement(schema._def.innerType),
            });
        } else if (schema instanceof z.ZodLazy) {
            return new z.ZodLazy({
                ...schema._def,
                // NB: This leaks `fn` into the schema, but there is no other way to support recursive schemas
                getter: () => mapElement(schema._def.getter()),
            });
        } else if (schema instanceof z.ZodBranded) {
            return new z.ZodBranded({
                ...schema._def,
                type: mapElement(schema._def.type),
            });
        } else if (schema instanceof z.ZodEffects) {
            return new z.ZodEffects({
                ...schema._def,
                schema: mapElement(schema._def.schema),
            });
        } else if (schema instanceof z.ZodFunction) {
            return new z.ZodFunction({
                ...schema._def,
                args: schema._def.args.map((arg: z.ZodTypeAny) => mapElement(arg)),
                returns: mapElement(schema._def.returns),
            });
        } else if (schema instanceof z.ZodPromise) {
            return new z.ZodPromise({
                ...schema._def,
                type: mapElement(schema._def.type),
            });
        } else if (schema instanceof z.ZodCatch) {
            return new z.ZodCatch({
                ...schema._def,
                innerType: mapElement(schema._def.innerType),
            });
        } else if (schema instanceof z.ZodTuple) {
            return new z.ZodTuple({
                ...schema._def,
                items: schema._def.items.map((item: z.ZodTypeAny) => mapElement(item)),
                rest: schema._def.rest && mapElement(schema._def.rest),
            });
        } else if (schema instanceof z.ZodDiscriminatedUnion) {
            const optionsMap = new Map(
                [...schema.optionsMap.entries()].map(([k, v]) => [
                    k,
                    mapElement(v) as ZodDiscriminatedUnionOption<string>,
                ]),
            );

            return new z.ZodDiscriminatedUnion({
                ...schema._def,
                options: [...optionsMap.values()],
                optionsMap: optionsMap,
            });
        } else if (schema instanceof z.ZodUnion) {
            return new z.ZodUnion({
                ...schema._def,
                options: schema._def.options.map((option: z.ZodTypeAny) => mapElement(option)),
            });
        } else if (schema instanceof z.ZodIntersection) {
            return new z.ZodIntersection({
                ...schema._def,
                right: mapElement(schema._def.right),
                left: mapElement(schema._def.left),
            });
        } else if (schema instanceof z.ZodRecord) {
            return new z.ZodRecord({
                ...schema._def,
                keyType: mapElement(schema._def.keyType),
                valueType: mapElement(schema._def.valueType),
            });
        } else {
            return schema;
        }
    }

    return fn(mapInner());
}

export function deepPartial<T extends z.ZodTypeAny>(schema: T): T {
    return mapOnSchema(schema, (s) => (s instanceof z.ZodObject ? s.partial() : s)) as T;
}

/** Make all object schemas "strict" (ie. fail on unknown keys), except if they are marked as `.passthrough()` */
export function deepStrict<T extends z.ZodTypeAny>(schema: T): T {
    return mapOnSchema(schema, (s) =>
        s instanceof z.ZodObject && s._def.unknownKeys !== "passthrough" ? s.strict() : s,
    ) as T;
}

export function deepStrictAll<T extends z.ZodTypeAny>(schema: T): T {
    return mapOnSchema(schema, (s) => (s instanceof z.ZodObject ? s.strict() : s)) as T;
}

This is missing ZodDefault. Is that not necessary? I noticed this when I created another implementation to dynamically (and with async) walk and modify the schema where I had someOtherObjectSchema.array().default([]).

My solution adds this:

    } else if (schema instanceof z.ZodDefault) {
      const innerType = mapElement(schema.removeDefault());
      return z.ZodDefault.create(innerType, { default: schema._def.defaultValue });

@kernwig That's interesting...

The type list is from the official Zod source code, so if you needed that, it implies that the deepPartialify implementation in Zod itself is buggy.

Indeed, analyzing the rest of the source code of Zod, this (and deepPartialify) is also missing ZodLazy and ZodEffects and a bunch of other stuff.

@jaens I did indeed need it for an applications where I dynamically add a superRefine to some string properties that may have a default value, or a string property on an object inside of an array that defaults to empty.

Aside: Because my callback fn needs to call an API to fetch the list of valid string values, I had to implement an async version of mapOnSchema.

Just discovered that this function is why my object properties loose their z.describe() values when I use zod-to-openapi to create OpenAPI.

I fixed by adding { ...schema._def } as the second parameter to all of the z.Zod<thing>.create() function calls, thus preserving not only the description but errorMap, invalid_type_error, and required_error.

I updated the code to account for all known flaws. The usage of .create() was removed to be more "extension-proof".

Aside: Because my callback fn needs to call an API to fetch the list of valid string values, I had to implement an async version of mapOnSchema.

I mean, technically might not be necessary, it can be done in three passes (since the mapping order is deterministic, unless using lazy schemas):

1. Use mapOnSchema with the identity function, but as a side-effect, push all required requests into an eg. array (as promises).
2. Await all the promises in the array, into a new resolved one.
3. Run mapOnSchema again, this time reading the resolved values from the array (with an incrementing index).

Updated the solution to with more details and a license.

@jaens Thanks for the update. Will be great if you can add JS doc comments to all exported functions to clearly differentiate their usage.

Looks time to turn this gist into a published library.

Good job up there!

However, I encounter a bug with array of union objects. Works only first type in union:

The test

import { z } from "zod";

import { deepPartial } from "./zodDeep.utils";

const schema = z.object({
  arr: z.array(
    z.union([
      z.object({
        union1_Key1: z.string(),
        union1_Key2: z.number(),
      }),
      z.object({
        union2_Key1: z.string(),
        union2_Key2: z.number(),
      }),
    ]),
  ),
});

const data = {
  arr: [
    // {
    //   union1_Key1: "d",
    //   union1_Key2: 2,
    // },
    {
      union2_Key1: "f",
      union2_Key2: 3,
    },
  ],
};

describe("zodDeep.utils", () => {
  it("shoud make schema optional and result in sucess parse", () => {
    const _schema = deepPartial(schema);
    // const _schema = schema;

    expect(_schema.parse(data)).toEqual(data);
  });
});

Test result

---

UPD: just realized that simple union objects doesn't work either

Simple union test

const union1Schema = z.union([
  z.object({ union1_Key1: z.string(), union1_Key2: z.number() }),
  z.object({ union2_Key1: z.string(), union2_Key2: z.number() }),
]);

const unionData = { union2_Key1: "f", union2_Key2: 3 };

it("shoud make schema optional and result in sucess parse", () => {
  const _schema = deepPartial(union1Schema);
  // const _schema = union1Schema;

  expect(_schema.parse(unionData)).toEqual(unionData);
});

@Alexander-ilyin3 No, this library does work perfectly (ie. outputs the "correct" schema) in this case.

You just missed the fact that this is in fact how Zod works:

const union1SchemaPartial = z.union([
    z.object({ union1_Key1: z.string(), union1_Key2: z.number() }).partial(), // <<<<< NB
    z.object({ union2_Key1: z.string(), union2_Key2: z.number() }).partial(), // <<<<< NB
]);

const unionData = { union2_Key1: "f", union2_Key2: 3 };

test("should make schema optional and result in sucess parse", () => {
    // Nope, this is not how Zod works
    expect(union1SchemaPartial.parse(unionData)).toEqual(unionData);
});

To explain:

Zod will choose the first successful parse from an union. Since an empty object is a successful parse of an object with all keys optional (and also type-correct), the empty object is a perfectly "valid" result. (although unintuitive)

For this and many other reasons, in well-designed APIs, all unions should have at least 1 non-optional field, eg. a discriminator.
(or have completely non-overlapping types, eg. a string and an object)

If you want to union objects, do it via .merge() or .extend() etc. instead.

Ok. I guess I using it the wrong way then.

What I was trying to do is to make the first parse regular, then if that catches an error try to do deepPartial on a schema and then try to parse it that way.
So: try return schema.parse(data) catch try return deepPartial(schema).parse(data) catch return data.

The purpose was to not throw an error when some of the fields are missing because of the unstable backend. Just for application to try to work with DTOs "as is" and at the same time to evaluate .transform() on schemas that uses it. And if that fails, well, you on your own with the raw data.

Vanilla zod will not do a transform on a second schema in union if first passes the validation. Wow. That is so unfortunate... But it perfectly makes sense for me now

I realized that I need to write another type of Frankenstein in my case. And also became a better developer

I've given it a go, but I don't think I understand it enough to translate this into working with zod v4. Is it possible to produce one which works, or give me some hints as to what needs to change?

Hey @jaens, I tagged you in the deepPartial thread of the zod repo, but wanted to follow up here as well, since I think you might find the project interesting.

I've been experimenting with getting recursion schemes working in TypeScript, and found a pretty nice (and type-safe) way to do it :)

I credited you in the source code for the Functor implementation, since circular schema detection is an edge case I hadn't thought of until I stumbled on your gist.

I've implemented a handful of schemes, the one I use most is the catamorphism (my implementation is technically a paramorphism). A bunch of scary names, but once you pick up one, the rest come pretty naturally.

There are a few others I've implemented but haven't found occasion to use yet.

If any of this sounds interesting to you, I'd be down to collaborate on something! Using these feels very magical.

Also I've been experimenting with different ways of introducing recursion schemes to a broader TypeScript audience.

My first attempt at making them more accessible was a blogpost, but I'm not happy with it, maybe you'll have some ideas.

I know @josephjunker has written about them before. Recently we've discussed possibly working on something together. If you're interested in collaborating, the more the merrier I'm sure :)

@ahrjarrett Hah ya crazy lad, looks like you did it... I off-handedly mentioned Functors in my comments, although I did have F-algebras and HKT emulation etc. in the back of my mind, the margin wasn't big enough to write that note ;)

I think transforming Zod (or schemas in general) is probably the closest one would get to recursion schemes and all that jazz in "everyday" TypeScript programming (...very likely outnumbering people eg. implementing programming languages or DSLs). You kind of have to have a "wide" recursive sum type (like a Zod schema) to actually feel the need for it, otherwise just "bruteforcing" (ie. inlining) the recursion for every transformer is the "obvious" option.

(general thoughts follow)

Another theoretical concept that I feel is underutilized in TS is lenses... every time I work on a complicated form in React I just facepalm at the fragile boilerplate required to keep everything coherent manually.

I think the current TS FP libraries suffer a bit from trying to be the "one true way", having absolutely every FP abstraction under the sun, so for someone starting out the eyes glaze over...

...and as much as I appreciate FP abstractions for their clarity in guiding thought and generally "railroading" you into reliable and complete solutions, eg. I would never use a standard/trivial Result<> type in production programming, as the debugging and other "real engineering" properties of it are [censored]. (take a dive into the world of Rust advanced error handling crates to see the minimum what of I would consider an acceptable solution over exceptions)

---

and yeah, I certainly have wanted for these ideas to reach a larger audience...

Another theoretical concept that I feel is underutilized in TS is lenses... every time I work on a complicated form in React I just facepalm at the fragile boilerplate required to keep everything coherent manually.

Couldn't agree more. You might be interested zx.makeLens, which is something I've been meaning to get back to working on. Users can pass a zod schema and a "selector" function (which is a just a Proxy that tracks property access) to get back a lens, prism or traversal (depending on what they selected).

Under the hood, the implementation uses the Profunctor encoding, but I've also been meaning to try out the existential encoding.

Support is still experimental, it's mostly just a POC for now.