Approach for handling `discriminatedUnion` with `superRefine` with Zod

Question

Dynamic Validation Schema

1. Base Schema: Create a base schema that contains common fields.

2. Dynamic Validations:

• Code: Implement dynamic validations based on the code field.
• Flag: Use a flag passed as a parameter to determine if additional validations should be applied.

3. Shared complex validations: Use superRefine in the base schema to maintain shared validation logic across all schemas.

Code example

type Codes = {
    One: 'one',
    Two: 'two',
}

const BaseSchema = z.object({
    code,
    // ... other common fields
}).superRefine((data, ctx) => {
    // Common complex validation logic
});

const One = z.object({
  email,
});
// ... more specific schemas like `Two`

const FlaggedSchema = z.object({
  dateOfBirth,
});

const createSchema = (code: Codes, isFlagged: boolean) => {
    const isOne = code === Codes.One
    const baseSchema = BaseSchema.extend({
        code: z.literal(code),
        ...(isOne ? One.shape : {}),
        ...(isFlagged ? FlaggedSchema.shape : {}),
    });

    return baseSchema;
};

export const createMapSchema = (isFlagged: boolean) => {
  const schemas = Object.values(Codes).map((code) =>
    createSchema(code, isFlagged),
  );
  const unions = z.discriminatedUnion('code', schemas);
  const map = z.record(z.array(unions));
  const schema = z.object({
    types: map,
  });
  return schema;
};

Problem

superRefine it cannot be applied to the base schema if you plan to use schema methods like merge afterward. This limitation arises because superRefine returns a Zod effect, which prevents further modifications to the schema.

so, how to apply common complex validations(refine or superRefine) in the base schema but also running dynamic validations depending schema fields or external data?

Aditional notes

I reviewed this question but I didn't see anything to use with this specific case

Answer 1

Instances of the class ZodEffects (as well as the other ZodType descendants) have a _def property that contains all the required information to reconstitute the object from its components. See the source code types.ts#L89.

A ZodEffects object zodEffects that is generated by .refine or .superRefine has a zodEffects._def.schema which points to a clone of the original ZodType to which .refine or .superRefine was applied and a zodEffects._def.effect.refinement that contains the refinement method arguments.

The following (experimental) function applyToRefined applies methods like extend to zodType that can be either a ZodObject or a ZodEffects obtained by refining a ZodObject.

It follows the chain of _defs until it finds a non-refined object, applies to that object the method (e.g., extend) and then reapplies the refinements:

function applyToRefined(zodType: z.ZodType, method: string, ...args: any[]): z.ZodType{
    const refinements: any[] = [];
    while(zodType instanceof z.ZodEffects){
        if(zodType._def?.effect?.type !== 'refinement'){
            throw new Error('Cannot handle object that is not produced by .refine or .superRefine');
        }
        refinements.push(zodType._def.effect.refinement);
        zodType = zodType._def.schema;
    }

    if(!zodType[method] || !zodType[method].apply){
        throw new Error(`The schema of the ZodEffects object doesn't have the method ${method}`);
    }
    zodType = zodType[method].apply(zodType, args);
    while(refinements.length > 0){
        const refinement = refinements.shift();
        zodType = zodType.superRefine(refinement);
    }
    return zodType;
}

A test case, based on an example from the docs (since the example in the question was too fragmentary for me to use, but the scope should be the same):

const BaseSchema = z.object({
        first: z.string(),
        second: z.number(),
    })
    .superRefine((arg: any, ctx: any) => {
        if(Math.round(arg.second) !== arg.second){
            ctx.addIssue({
                code: z.ZodIssueCode.invalid_type,
                message: ".second should be integer",
            });
        }
    })
    .refine((arg: any) => arg.first.length > 0, {
        message: ".first should not be the empty string"
    });

function dynamicExtend(zodType: z.ZodType, code: string, otherStringProp: string){
    return applyToRefined(zodType, 'extend', {
        code: z.literal(code),
        [otherStringProp]: z.string()
    });
}

const BaseSchema1 = dynamicExtend(BaseSchema, 'one', 'andOne');
const BaseSchema2 = dynamicExtend(BaseSchema, 'two', 'andTwo');

console.log(BaseSchema1.parse({first: 'one', second: 100, code: 'one', andOne: '1'}));
console.log(BaseSchema2.parse({first: 'one', second: 100, code: 'two', andTwo: '2'}));
try{
    console.log(BaseSchema1.parse({first: '', second: 100.1, code: 'one', andOne: '2'}));
}
catch(err){
    console.error(err);
}
try{
    console.log(BaseSchema1.parse({first: '', second: 100.1, code: 'two'}));
    //see https://github.com/colinhacks/zod/discussions/2971#discussioncomment-7588935
}
catch(err){
    console.error(err);
}

In order to create a discriminated union of ZodEffects, we can apply the same strategy as in applyToRefined, but to an array of ZodEffects:

function combineRefined(
    zodTypes: z.ZodType[],
    combinatorFunction: (...args: z.ZodType[]) => z.ZodType
): z.ZodType {
    const refinements: any[] = [];
    const zodSchemas: z.ZodType[] = [];
    let zodType: z.ZodType;
    for(zodType of zodTypes){
        while(zodType instanceof z.ZodEffects){
            if(zodType._def?.effect?.type !== 'refinement'){
                throw new Error('Cannot handle object that is not produced by .refine or .superRefine');
            }
            refinements.push(zodType._def.effect.refinement);
            zodType = zodType._def.schema;
        }
        zodSchemas.push(zodType);
    }
    zodType = combinatorFunction(...zodSchemas);
    while(refinements.length > 0){
        const refinement = refinements.shift();
        zodType = zodType.superRefine(refinement);
    }
    return zodType;
}


// create the union of the BaseSchema1, BaseSchema2 above, discriminated by "code" property
const BaseSchema12 = combineRefined([BaseSchema1, BaseSchema2],
    (...schemas: z.ZodType[]) => z.discriminatedUnion("code", schemas));

And testing the schema with .parse:

console.log('--------------------');
console.log(BaseSchema12.parse({code: 'one', first: 'one', second: 100, andOne: '1'}));
console.log(BaseSchema12.parse({code: 'two', first: 'one', second: 100, andTwo: '2'}));
try{
    console.log(BaseSchema12.parse({code: 'one', first: 'one', second: 100, andTwo: '1'}));
    // error: no andOne
}
catch(err){
    console.error(err);
}

try{
    console.log(BaseSchema12.parse({code: 'one', first: '', second: 100.1, andOne: '1'}));
    // refine errors for first and second
}
catch(err){
    console.error(err);
}

The whole typescript file in a stackblitz (results in the browser console).

One could consider, for convenience, adding such methods as extend or discriminatedUnion directly to the ZodEffects class, but that's probably too hacky for an already hacky solution.