Deprecating `z.discriminatedUnion`?

[colinhacks]

Superceded by #3407

---

I'm planning to deprecate z.discriminatedUnion in favor of a "switch" API that's cleaner and more generalizable. You can dynamically "switch" between multiple schemas at parse-time based on the input.

const schema = z.switch(()=>{
  return Math.random() ? z.string() : z.number()
});

I expand more on the z.switch API later. Let's talk about z.discriminatedUnion.

Why

z.union naively tries each union element until parsing succeeds. That's slow and bad. Zod needed some solution.

z.discriminatedUnion was a mistake. The API was good but I had reservations about the implementation. It required fiddly recursive logic to exrtract a literal discriminator key from each union element.

It's a bad sign when a method or class requires weird recursive traversal of other schemas. For starters, Zod is designed to be subclassable. Users can theoretically subclass ZodType to implement custom schema types. But logic like this instanceof switch statement don't and can't account for any user-land schema types.

But the main problem is just that this kind of pattern is bad and introduces a lot of edge cases. It means that only certain kinds of schemas are allowed as discriminators, and others will fail in unexpected ways. There are now dozens of issues that have been opened regarding these various edge cases. The PRs attempting to solve this problem are irredeemably complex and introduce even more edge cases.

• Enumerated discriminator #1773
• [BETA 3.20] Using transforms with discriminatedUnion #1667
• discriminatedUnion with recursive type results in ts error #1384
• Invalid type error with discriminatedUnion and or #1757
• Is there a way to combine discriminatedUnion with lazy? #1504
• Relationship between schema object fields #1136
• discriminatedUnion errors with z.object().transform() #1477
• Transformed objects in .union() vs .discriminatedUnion() #1424
• Zod discriminated union discriminator value accepts zod enum #1213
• discriminatedUnion produces TS error when .default or .preprocess are applied #1490
• Type error with discriminatedUnion when using dynamic amount of objects #1181
• Allow more types in discriminated unions #1075
• feature request: combining/nesting/extending discriminated unions #1444
• Discriminated Union based on nested object property? #1868
• discriminatedUnion with array element as discriminator #1878

The .deepPartial API has this same problem. I'm deprecating it for the same reason.

Many of those issues are asking for non-literal discriminator types:

type MyUnion = 
  | { type: "a", value: string }
  | { type: "b", value: string }
  | { type: "c", value: string }
  | { type: string, value: string }
  | { type: null, value: string }
  | { type: undefined, value: string }
  | { type: MyEnum, value: string }
  | { type: { nested: string }, value: string }
  | { type: number[], value: string };

Imagine each of those elements are represented with Zod schemas. Zod would need to extract the type field from each of these elements and find a way to match the incoming input.type against those options. In the general case, Zod would extract the type field from the shape of each component ZodObject and check input.type against those schemas until a match is found. At that point, we're back to doing a parse operation for each element of the union, which is what z.discriminatedUnion is supposed to avoid doing. (It's still doing less work than the naive z.union but still.)

Another issue is composability. The existing API expects the second argument to be an array of ZodObject schemas.

z.discriminatedUnion("type", [
  z.object({ type: z.literal("a"), value: z.string() }),
  z.object({ type: z.literal("b"), value: z.string() }),
])

This isn't composable, in that you can't nest discriminated unions or add additional members.

const ab = z.discriminatedUnion("type", [
  z.object({ type: z.literal("a"), value: z.string() }),
  z.object({ type: z.literal("b"), value: z.string() }),
]);

const abc = z.discriminatedUnion("type", [
  ab,
  z.object({ type: z.literal("c"), value: z.string() }),
]);

• Nested discriminated unions #1618
• Nested discriminated unions #1610
• discriminated union cant be nested #1884

Yes, Zod could support both (ZodObject | ZodDiscriminatedUnion)[] as union members, but that requires additional messy logic that reflects a more fundamental problem with the API. It also makes increasingly difficult to enforce typesafety on the union - it's important that all union elements have a type property, otherwise the union is no longer discriminable.

Replacement: z.switch

const schema = z.switch(input => {
  return (typeof input) === "string" ? z.string() : z.number();
})

schema.parse("whatever"); // string | number

A discriminated union looks like this:

const schema = z.switch((input) => {
  switch(input.key){
    case "a":
      return z.object({ key: z.literal("a"), value: z.string() })
    case "b":
      return z.object({ key: z.literal("b"), value: z.number() })
    default:
      return z.never()
  }
});
schema.parse({ /* data */ });
// { key: 'a', value: 'asdf' } | { key: 'b', value: number }

Ultimately the z.switch API is a far more explicit and generalizable API. Zod doesn't do any special handling. The user specifies exactly how the input will be used to select the schema. z.switch() accepts a function. The ResultType of that function is inferred. It will be the union of the schema types returned along all code paths in the function. For instance:

const schema = z.switch(()=>{
  return Math.random() ? z.string() : z.number()
});

Zod sees that the return type of the switcher function is ZodString | ZodNumber. The result of the z.switch is ZodSwitch<ZodString | ZodNumber>. The result of schema.parse(...) is string | number.

You can represent discriminated unions explicitly like this:

const schema = z.switch((input) => {
  switch(input.key){
    case "a":
      return z.object({ key: z.literal("a"), value: z.string() })
    case "b":
      return z.object({ key: z.literal("b"), value: z.number() })
    default:
      return z.never()
  }
});
schema.parse({ /* data */ });
// { key: 'a', value: 'asdf' } | { key: 'b', value: number }

This can be written in a more condensed form like so:

const schema = z.switch((input) => ({
  a: z.object({ key: z.literal("a"), value: z.string() }),
  b: z.object({ key: z.literal("b"), value: z.number() }),
}[input.key as string]));
  
schema.parse({ key: 'a', value: 'asdf' });
// { key: 'a', value: 'asdf' } | { key: 'b', value: number }

It's marginally more verbose. It's also explicit, closes 30+ issues, eliminates a lot of hairy logic, and lets Zod represent the full scope of TypeScript's type system. z.discrimininatedUnion is too fragile and causes too much confusion so it needs to go.

[mgreenw]

While I really liked the simple API of discriminatedUnion and am sad to see it go, this makes sense. Thank you for the detailed write up and your work to make a replacement that’s better!

[dvargasfuertes]

You can represent discriminated unions explicitly like this

Any reason why discriminatedUnion can't now call the new switch method in the way you described as its new implementation?

Or, introduce a new API that just accepts an array of schemas that then calls the switch function under the hood.

One strength of zod is how closely defining schemas resembles defining the typescript types. Requiring a function to be defined in userland for unions deviates from how that user would define the discriminated union in standard typescript.

To solve the recursive issue, there could be a requirement that only allows top level discriminators, or that the discriminator defined would need to define the full path to the key in a lodashy way, like key.subpath.

[JacobWeisenburger]

@colinhacks thanks for making a better system! Can't wait to try it out. Do you have an idea of when it will be released?

[dgritsko]

Currently, z.infer plays really nicely with z.discriminatedUnion. For example, given the following:

const ab = z.discriminatedUnion("type", [
  z.object({ type: z.literal("a"), value: z.string() }),
  z.object({ type: z.literal("b"), value: z.string() }),
]);

type AB = z.infer<typeof ab>

This will result in AB being defined as:

type AB = {
  type: "a";
  value: string;
} | {
  type: "b";
  value: string;
}

Which is exactly what I would expect. How will type inference work with z.switch -- will it support inference of "idiomatic" type definitions, similar to this example?

[pato1]

I really like the new API as it opens up new options for building any validation logic.
A thing to consider is the type of input for the switch function. In my opinion, you would almost always need to add switch after some preliminary validation.
Given the example replacing distributed union you've provided:

const schema = z.switch((input) => {
  switch(input.key){
    case "a":
      return z.object({ key: z.literal("a"), value: z.string() })
    case "b":
      return z.object({ key: z.literal("b"), value: z.number() })
    default:
      return z.never()
  }
});
schema.parse({ /* data */ });

you would probably want to validate data is an object with property key and some additional properties:

z.object({key: z.enum(['a', 'b'] as const)}).passthrough();

And then you should have the input parameter of the switch function typed accordingly (input: {key: 'a' | 'b'}).

Therefore I would suggest implementing switch similar to ZodPipeline in that it can be added to existing validation.
Then it could be used as follows:

const schema = z.object({key: z.enum(['a', 'b'] as const)}).passthrough().switch((input) => {
  // Here `input` is of type `{key: 'a' | 'b'}`
  switch(input.key){
    case "a":
      return z.object({key: z.literal("a"), value: z.string()})
    case "b":
      return z.object({key: z.literal("b"), value: z.number()})
    default:
      return z.never()
  }
});
schema.parse({ /* data */ });

Inability to compose distributed unions led me to build a temporary switch type based on the implementation of ZodPipeline:

export interface ZodSwitchDef<T extends ZodTypeAny, B extends ZodTypeAny = ZodType<unknown>> extends z.ZodTypeDef {
  getSchema: (value: B['_output']) => T
  base: B
  typeName: 'ZodSwitch'
}
export class ZodSwitch<
  T extends ZodTypeAny,
  B extends ZodTypeAny
  > extends ZodType<T['_output'], ZodSwitchDef<T, B>, B['_input']> {
  _parse(input: z.ParseInput): z.ParseReturnType<any> {
    const {ctx} = this._processInputParams(input)
    if (ctx.common.async) {
      const handleAsync = async () => {
        const inResult = await this._def.base._parseAsync({
          data: ctx.data,
          path: ctx.path,
          parent: ctx,
        })
        if (inResult.status !== 'valid') return z.INVALID
        return this._def.getSchema(inResult.value)._parseAsync({
          data: inResult.value,
          path: ctx.path,
          parent: ctx,
        })
      }
      return handleAsync()
    } else {
      const inResult = this._def.base._parseSync({
        data: ctx.data,
        path: ctx.path,
        parent: ctx,
      })
      if (inResult.status !== 'valid') return z.INVALID
      return this._def.getSchema(inResult.value)._parseSync({
        data: inResult.value,
        path: ctx.path,
        parent: ctx,
      })
    }
  }

  static create<T extends ZodTypeAny, B extends ZodTypeAny>(
    getSchema: (value: B['_output']) => T,
    base: B,
  ): ZodSwitch<T, B> {
    return new ZodSwitch({
      base,
      getSchema,
      typeName: 'ZodSwitch',
    })
  }
}

// ...

class ZodType {
 switch<T extends ZodTypeAny>(getSchema: (value: Output) => T): ZodSwitch<T, this> {
   return ZodSwitch.create(getSchema, this);
 }
}

Note, unlike ZodPipeline, if "base" validation failed, I had to return INVALID, as requirement for the switch function (getScheme) parameter is not fulfilled.

Inference works as expected:

const ab = ZodSwitch.create(({type}) => {
  switch (type) {
    case 'a':
      return z.object({type: z.literal('a'), value: z.string()})
    case 'b':
      return z.object({type: z.literal('b'), value: z.number()})
    default:
      return z.never()
  }
}, z.object({type: z.enum(['a', 'b'] as const)}).passthrough())

type AB = z.infer<typeof ab>

// type AB is {type: "a", value: string} | {type: "b", value: number}

Is this somewhat related to what you have in mind for the new API?

[colinhacks]

Is this somewhat related to what you have in mind for the new API?

@pato1 Almost exactly, aside from base. I like the idea of having some "pre-validation" so you get some type safety before attempting the switch logic, and a .switch() method is certainly interesting. I'd propose to go a step further and overload the .pipe() API to accept (input: unknown)=>ZodType. They're conceptually the same - switch is a dynamic pipe.

// side node: you don't need `as const` in `z.enum`
z.object({key: z.enum(['a', 'b'])}).pipe((input) => {
  switch(input.key){
    case "a":
      return z.object({key: z.literal("a"), value: z.string()})
    case "b":
      return z.object({key: z.literal("b"), value: z.number()})
    default:
      return z.never()
  }
});

As with your ZodSwitch class, we just look at the return type of the "switcher function", which will be some union of all the schema types that are returned along the various code paths. Inference will work as expected. cc @dgritsko

[colinhacks]

Unfortunately though, over the course of trying to explain the problems with z.discriminatedUnion I hit on an implementation that actually solves most of my issues with it.

In the general case, Zod would extract the type field from the shape of each component ZodObject and check input.type against those schemas until a match is found. At that point, we're back to doing a parse operation for each element of the union, which is what z.discriminatedUnion is supposed to avoid doing. (It's still doing less work than the naive z.union but still.)

Conceptually, we could loop over the set of options (ZodObject[]) in the discriminated union, pull out <option>.shape[discriminator], and use that to parse input[discriminator]. The first time this succeeds, we have a match, and use that option to validate the whole input. I was originally presenting this as a straw man for why ZodDiscriminatedUnion is inherently flawed...but actually it's a pretty good approach. With this, there's no need to spelunk into the internals of the union's options to extract out literal discriminator values. There will be a slight performance tradeoff relative to the current implementation since we're still executing a .parse operation for each option in the discriminated union. But most discriminator schemas should be simple & fast.

Some special logic is required to make nesting work (passing other ZodDiscriminatedUnion instances into z.discriminatedUnion), which is slightly annoying. It still may come back to bite me if people start asking to include ZodUnion, ZodEffects, ZodOptional, etc. as options in z.discriminatedUnion, at which point we're right back where we started, with a bunch of fiddly special-case logic.

[colinhacks]

Any reason why discriminatedUnion can't now call the new switch method in the way you described as its new implementation?

Because the hard part is extracting out the literal discriminator values from the schemas passed in ZodDiscriminatedUnion. We could rewrite z.discriminatedUnion to return a ZodSwitch but that would be a breaking change - better to just leave it as is and let people switch over at their liesure.

Or, introduce a new API that just accepts an array of schemas that then calls the switch function under the hood.

How is this different from the current API for z.discriminatedUnion? The problem is that it's hard and messy.

One strength of zod is how closely defining schemas resembles defining the typescript types. Requiring a function to be defined in userland for unions deviates from how that user would define the discriminated union in standard typescript.

Remember that there is no syntax for "discriminated unions" in typescript. There are just unions. Zod only includes a separate API as a performance optimization when parsing unions, and I think z.switch actually makes this more explicit.

[dvargasfuertes]

How is this different from the current API for z.discriminatedUnion?

I realized here what I was proposing was actually the API of z.union, which runs into the problem you described in your description.

[pato1]

  I'd propose to go a step further and overload the .pipe() API to accept (input: unknown)=>ZodType. They're conceptually the same - switch is a dynamic pipe.

@colinhacks I like this idea, even more now you're considering rewriting z.discriminatedUnion. But shouldn't then the "switcher function" accept (input: A["_output"])=>ZodType?

I hit on an implementation that actually solves most of my issues with it.

Also one benefit to making z.discriminatedUnion work, is generators (Zod to X) like @anatine/zod-openapi.
With the switch API there would be no way of getting individual options for schema generation.

Conceptually, we could loop over the set of options (ZodObject[]) in the discriminated union, pull out .shape[discriminator], and use that to parse input[discriminator].

This is just an idea, bear with me. One other way to achieve that even for other types of Options that closely resembles TypeScript would be to provide "indexed access". If types had a method, let's say index(key: string): ZodType than you could just use that instead of <option>.shape[discriminator]. Every type that implemented this method (e.g. ZodDiscriminatedUnion) could then be used in discriminated union option. I can see it being useful even in other cases, for example if I want to validate only one field against some complex schema.

Yes, it would be calling possibly deep for some complex schemas, but it would use public API of those types, instead of poking around in internals.

Pseudo examples:

ZodObject.index = (key) => this.shape[key]
ZodIntersection.index = (key) => z.intersection(this.left.index(key), this.right.index(key)
ZodUnion.index = (key) => z.union(this.options.map((option) => option.index(key))

[sluukkonen]

FWIW, the issue can also be solved with an API like this

z.discriminatedUnion("type", [
  ["a", z.object({ type: z.literal("a"), value: z.string() })],
  ["b", z.object({ type: z.literal("b"), value: z.string() })],
])

It's a bit more verbose, but the implementation doesn't have to assume anything about the structure of the inner types.