How to omit/pick a value from an ZodEnum?

[PhilBookst]

Let's say I have an enum with different medical occupations:

export const OccupationSchema = z.enum([
    "doctor",
    "pharmacist",
    "nutritionist",
    "occupational-therapist",
])

and a faculty schema:

export const FacultySchema = z.enum([
    "general-medicine",
    "anesthesiology",
    "anatomy",
    "occupational-medicine",
    "ophthalmology",
])

I need to create a discreminatedUnion schema that checks if the occupation is === "doctor" to allow a faculty to be set:

const MedicOccupationDiscriminator = z.discriminatedUnion("occupation", [
    z.object({
        occupation:  // something like z.pick(OccupationSchema, "doctor"),
        faculty: FacultySchema,
    }),
    z.object({
        occupation: // something like z.omit(OccupationSchema, "doctor"),
    }),
]);

How can I achieve this behaviour? Thank you!

[JacobWeisenburger]

Is this what you are looking for?

export const OccupationSchema = z.enum( [
    'doctor',
    'pharmacist',
    'nutritionist',
    'occupational-therapist',
] )

type OccupationOptions = typeof OccupationSchema[ 'options' ]

export const FacultySchema = z.enum( [
    'general-medicine',
    'anesthesiology',
    'anatomy',
    'occupational-medicine',
    'ophthalmology',
] )

type ExcludeFromTuple<T extends readonly any[], E> =
    T extends [ infer F, ...infer R ] ? [ F ] extends [ E ] ? ExcludeFromTuple<R, E> :
    [ F, ...ExcludeFromTuple<R, E> ] : []

const MedicOccupationDiscriminator = z.discriminatedUnion( 'occupation', [
    z.object( {
        occupation: z.literal( OccupationSchema.enum.doctor ),
        faculty: FacultySchema,
    } ),
    z.object( {
        occupation: z.enum(
            OccupationSchema.options
                .filter( x => x !== OccupationSchema.enum.doctor ) as ExcludeFromTuple<
                    OccupationOptions,
                    typeof OccupationSchema.enum.doctor
                >
        ),
    } ),
] )

console.log(
    MedicOccupationDiscriminator.safeParse( {
        occupation: 'doctor',
        faculty: 'anesthesiology',
    } ).success
) // true

console.log(
    MedicOccupationDiscriminator.safeParse( {
        occupation: 'pharmacist',
    } ).success
) // true

console.log(
    MedicOccupationDiscriminator.safeParse( {
        occupation: 'foo',
    } ).success
) // false

[PhilBookst]

Wow this is it, thank you!

Is there a way to extend a zod object with this MedicOccupationDiscriminator?

I have a object with this shape:

const PersonSchema = z.object({
    firstName: z.string(),
    lastName: z.string(),
    // I need the disriminator right here 
    // ...MedicOccupationDiscriminator this would be ideal, but sadly doesn't work 
});

If I try MedicSchema.and(MedicOccupationDiscriminator ) this works fine, but I still need the omit function on the schema to exclude certain fields. Is this possible or should I just create an extra schema without these fields?

Anyway thank you very much for this ts wizardry!

[JacobWeisenburger]

Is this what you are looking for?

const PersonSchema = z.object( {
    firstName: z.string(),
    lastName: z.string(),
} )

const MedicOccupationDiscriminator = z.discriminatedUnion( 'occupation', [
    PersonSchema.extend( {
        occupation: z.literal( OccupationSchema.enum.doctor ),
        faculty: FacultySchema,
    } ),
    PersonSchema.extend( {
        occupation: z.enum(
            OccupationSchema.options
                .filter( x => x !== OccupationSchema.enum.doctor ) as ExcludeFromTuple<
                    OccupationOptions,
                    typeof OccupationSchema.enum.doctor
                >
        ),
    } ),
] )

[PhilBookst]

Unfortunately not exactly, the omit function is sadly not available when designing the schema this way

Currently I have it like this:

const MedicSchema = z.discriminatedUnion("occupation", [
    PersonSchema.extend({
        occupation: z.literal(OccupationSchema.enum.doctor),
        faculty: FacultySchema,
    }),
    PersonSchema.extend({
        occupation: z.enum(
            OccupationSchema.options.filter(
                (x) => x !== OccupationSchema.enum.doctor
            ) as ExcludeFromTuple<
                OccupationOptions,
                typeof OccupationSchema.enum.doctor
            >
        ),
    }),
]);

but I'd like to have a CreateMedicSchema that for example excludes the email field like so:

const CreateMedicSchema = MedicSchema.omit({ email: true }) 

but after the discriminatedUnion omit sadly is not available. I don't know if this is currently possible as stated here.

const CreateMedicSchema = z.discriminatedUnion("occupation", [
    PersonSchema.extend({
        occupation: z.literal(OccupationSchema.enum.doctor),
        faculty: FacultySchema,
    }).omit({ email: true }),
    PersonSchema.extend({
        occupation: z.enum(
            OccupationSchema.options.filter(
                (x) => x !== OccupationSchema.enum.doctor
            ) as ExcludeFromTuple<
                OccupationOptions,
                typeof OccupationSchema.enum.doctor
            >
        ),
    }).omit({ email: true }),
]);

A bit ugly since there's a lot of duplicate code but well, maybe it's possible some day!

Anyway thank you very much for your help, much appreciated!

[JacobWeisenburger]

While working on something else I noticed .extract and remembered this discussion. Do you think it will help you?

const OccupationSchema = z.enum( [
    'doctor',
    'pharmacist',
    'nutritionist',
    'occupational-therapist',
] )

const FacultySchema = z.enum( [
    'general-medicine',
    'anesthesiology',
    'anatomy',
    'occupational-medicine',
    'ophthalmology',
] )

const MedicOccupationDiscriminator = z.discriminatedUnion( 'occupation', [
    z.object( {
        occupation: OccupationSchema.extract( [ 'doctor' ] ),
        faculty: FacultySchema,
    } ),
    z.object( {
        occupation: OccupationSchema.exclude( [ 'doctor' ] )
    } ),
] )

console.log(
    MedicOccupationDiscriminator.safeParse( {
        occupation: 'doctor', faculty: 'general-medicine'
    } ).success
) // true
console.log(
    MedicOccupationDiscriminator.safeParse( {
        occupation: 'pharmacist'
    } ).success
) // true

If you found my answer satisfactory, please consider supporting me. Even a small amount is greatly appreciated. Thanks friend! 🙏
https://github.com/sponsors/JacobWeisenburger