Support custom discriminators in `z.discriminatedUnion`

[RichardCPoint]

Currently z.discriminatedUnion only supports a restricted range of types for the discriminator field – mostly literal values.

There are use-cases for wanting to use other possible values (e.g. a union of values), or even custom types.

For example, an object might have different fields depending on the value of the type property, but we might also need a fallback case for any unknown value (i.e. any value other than the known values), e.g.:

enum Species {
  Cat = 'cat',
  Dog = 'dog',
}

const zPet = z.discriminatedUnion('species', [
  z.object({ species: z.literal(Species.Cat), lives: z.number() }),
  z.object({ species: z.literal(Species.Dog), breed: z.string() }),
  z.object({
    species: z.custom<string>(val => typeof val === 'string' && Object.values(Species).indexOf(val) < 0),
  }),
})

[ryami333]

A use-case I'd like to add is that sometimes the discriminator key is nested inside of a sub-object. For example, the Storyblok Rest API returns objects shaped like this:

type FooStory = {
  id: number,
  name: string,
  // … etc
  content: {
    component: "foo", // <-- Discriminated union key
    // … etc
  }
}

type BarStory = {
  id: number,
  name: string,
  // … etc
  content: {
    component: "bar", // <-- Discriminated union key
    // … etc
  }
}

So if I expected something to be FooStory | BarStory then I would hope to be able to write a schema like this:

z.discriminatedUnion("content.component", [
  z.object({
    id: z.number(),
    name: z.string(),
    // … etc
    content: z.object({
      component: z.literal("foo"),
      // … etc
    })
  }),
  z.object({
    id: z.number(),
    name: z.string(),
    // … etc
    content: z.object({
      component: z.literal("bar"),
      // … etc
    })
  }),
});

Seeing as Typescript now supports template literal types, I imagine this could be possible.