Is it possible to preprocess discriminator in zod discriminated union?

I need to coerce string formatted numbers to number before using a value in z.discriminatedUnion. Is this possible? Here is a simplified code snippet:

import { z } from "zod";

const BrushColorEnum = z.enum(
  ["BLUE_SILVER", "GREEN_SILVER", "RED", "GREEN_BLACK"],
  { message: "Invalid option" }
);

const BrushTypeEnum = z.enum(["THREAD", "MIXED", "CARLITE"], {
  message: "Invalid option",
});

const noBrushSchema = z.object({
  brush_qty: z.preprocess(
    (val) => (typeof val === "string" ? parseInt(val, 10) : val),
    z.literal(0)
  ),
  brush_type: z
    .union([
      z.string().length(0, { message: "Invalid option" }),
      z.undefined(),
      z.literal(false),
    ])
    .transform(() => undefined),
  brush_color: z
    .union([
      z.string().length(0, { message: "Invalid option" }),
      z.undefined(),
      z.literal(false),
    ])
    .transform(() => undefined),
});

const brushSchema = z.discriminatedUnion("brush_qty", [
  noBrushSchema,
  z.object({
    brush_qty: z.literal("2"),
    brush_type: BrushTypeEnum,
    brush_color: BrushColorEnum,
  }),
  z.object({
    brush_qty: z.literal("3"),
    brush_type: BrushTypeEnum,
    brush_color: BrushColorEnum,
  }),
]);

console.log(brushSchema.safeParse({ brush_qty: "0" }).error); // message: "Invalid discriminator value. Expected 0 | '2' | '3'"
console.log(brushSchema.safeParse({ brush_qty: 0 })); // success

Take a look at the brush_qty field. I would expect it to transform the string to number and I expect zod accept also the second safeParse, but it doesn't. It seems that the discriminator get validated before passing on to the schema part.

If in this case it is impossible to coerce it to number, what other choices than discriminatedUnion I have to get more or less the same functionality?

Thanks in advance!

Have brushSchema parse the relevant property first against a z.coerce.string() inside a z.object() with just that property. passthrough explicitly allows other properties, but I believe it is the default anyway -- just being clear about behaviour.

import { z } from 'zod';

const BrushColorEnum = z.enum(
  ['BLUE_SILVER', 'GREEN_SILVER', 'RED', 'GREEN_BLACK'],
  { message: 'Invalid option' }
);

const BrushTypeEnum = z.enum(['THREAD', 'MIXED', 'CARLITE'], {
  message: 'Invalid option',
});

const noBrushSchema = z.object({
  brush_qty: z.literal('0'),
  brush_type: z
    .union([
      z.string().length(0, { message: 'Invalid option' }),
      z.undefined(),
      z.literal(false),
    ])
    .transform(() => undefined),
  brush_color: z
    .union([
      z.string().length(0, { message: 'Invalid option' }),
      z.undefined(),
      z.literal(false),
    ])
    .transform(() => undefined),
});

const brushSchema = z
  .object({ brush_qty: z.coerce.string() })
  .passthrough()
  .pipe(
    z.discriminatedUnion('brush_qty', [
      noBrushSchema,
      z.object({
        brush_qty: z.literal('2'),
        brush_type: BrushTypeEnum,
        brush_color: BrushColorEnum,
      }),
      z.object({
        brush_qty: z.literal('3'),
        brush_type: BrushTypeEnum,
        brush_color: BrushColorEnum,
      }),
    ])
  );

console.log(
  brushSchema.safeParse({
    brush_qty: '2',
    brush_type: 'THREAD',
    brush_color: 'RED',
  })
);  // success

console.log(brushSchema.safeParse({ brush_qty: 0 })); // success

I don't believe you then need any of the other hacking in the other schema, but if you wanted to keep the ability to deal with floating points (e.g. make "2.2" match 2), you can use a similar trick:

const brushSchema = z
  .object({
    brush_qty: z.coerce.string().transform((arg) => parseInt(arg).toString()),
  })
  .passthrough()
  .pipe(
    z.discriminatedUnion('brush_qty', [

 // ... etc

This works by getting everything to a string first; then, if that string contains a floating point, that's normalised by doing a transform that parses it and returns to a string again.