Approach for handling discriminatedUnion with superRefine

[crisfcodes]

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

[nagyatka]

I use the following function when I have to put schemas with refines (ZodEffects) in a discriminated union.

function zDiscriminatedUnion<T extends readonly [z.ZodTypeAny, z.ZodTypeAny, ...z.ZodTypeAny[]]>(
    key: string,
    types: T,
): z.ZodUnion<T>
// eslint-disable-next-line @typescript-eslint/no-explicit-any
function zDiscriminatedUnion(key: string, types: z.ZodTypeAny[]): any {
    const optionsMap = new Map()
    for (const type of types) {
        const value = (type instanceof z.ZodEffects ? type.sourceType() : type).shape[key]
        if (!(value instanceof z.ZodLiteral) || optionsMap.has(value.value)) {
            throw new Error('cannot contruct discriminated union')
        }
        optionsMap.set(value.value, type)
    }
    return new z.ZodDiscriminatedUnion({
        typeName: z.ZodFirstPartyTypeKind.ZodDiscriminatedUnion,
        discriminator: key,
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        options: types as any,
        optionsMap,
    })
}

I hope it helps.

Luckily this issue will be solved in zod version 4: #2474 (comment)