Clean up discriminatedUnion

Author: Colin McDonnell

deno/lib/__tests__/discriminatedUnions.test.ts

@@ -15,11 +15,6 @@ test("valid", () => {
 });
 
 test("valid - discriminator value of various primitive types", () => {
-  interface lazy {
-    type: "lazy typed";
-    val: 17;
-  }
-
   const schema = z.discriminatedUnion("type", [
     z.object({ type: z.literal("1"), val: z.literal(1) }),
     z.object({ type: z.literal(1), val: z.literal(2) }),
@@ -30,43 +25,9 @@ test("valid - discriminator value of various primitive types", () => {
     z.object({ type: z.literal(null), val: z.literal(7) }),
     z.object({ type: z.literal("undefined"), val: z.literal(8) }),
     z.object({ type: z.literal(undefined), val: z.literal(9) }),
-    z
-      .object({ type: z.literal("transform"), val: z.literal(10) })
-      .transform((val) => ({
-        val: val.val,
-      })),
-    z
-      .object({ type: z.literal("refine"), val: z.literal(11) })
-      .refine(() => true),
-    z
-      .object({ type: z.literal("superRefine"), val: z.literal(12) })
-      .superRefine(() => {}),
-    z.lazy(() => z.object({ type: z.literal("lazy"), val: z.literal(13) })),
-    z.lazy(() =>
-      z
-        .object({ type: z.literal("chained 1"), val: z.literal(14) })
-        .transform((val) => ({
-          val: val.val,
-        }))
-    ),
-    z
-      .lazy(() =>
-        z.object({ type: z.literal("chained 2"), val: z.literal(15) })
-      )
-      .transform((val) => ({
-        val: val.val,
-      })),
-    z
-      .lazy(() =>
-        z.object({ type: z.literal("chained 3"), val: z.literal(16) })
-      )
-      .transform((val) => ({
-        val: val.val,
-      }))
-      .refine(() => true),
-    z.lazy(() =>
-      z.object({ type: z.literal("lazy typed"), val: z.literal(17) })
-    ) as z.ZodType<lazy>,
+    z.object({ type: z.literal("transform"), val: z.literal(10) }),
+    z.object({ type: z.literal("refine"), val: z.literal(11) }),
+    z.object({ type: z.literal("superRefine"), val: z.literal(12) }),
   ]);
 
   expect(schema.parse({ type: "1", val: 1 })).toEqual({ type: "1", val: 1 });
@@ -99,34 +60,6 @@ test("valid - discriminator value of various primitive types", () => {
     type: undefined,
     val: 9,
   });
-  expect(schema.parse({ type: "transform", val: 10 })).toEqual({
-    val: 10,
-  });
-  expect(schema.parse({ type: "refine", val: 11 })).toEqual({
-    type: "refine",
-    val: 11,
-  });
-  expect(schema.parse({ type: "superRefine", val: 12 })).toEqual({
-    type: "superRefine",
-    val: 12,
-  });
-  expect(schema.parse({ type: "lazy", val: 13 })).toEqual({
-    type: "lazy",
-    val: 13,
-  });
-  expect(schema.parse({ type: "chained 1", val: 14 })).toEqual({
-    val: 14,
-  });
-  expect(schema.parse({ type: "chained 2", val: 15 })).toEqual({
-    val: 15,
-  });
-  expect(schema.parse({ type: "chained 3", val: 16 })).toEqual({
-    val: 16,
-  });
-  expect(schema.parse({ type: "lazy typed", val: 17 })).toEqual({
-    type: "lazy typed",
-    val: 17,
-  });
 });
 
 test("invalid - null", () => {
@@ -196,9 +129,7 @@ test("wrong schema - missing discriminator", () => {
     ]);
     throw new Error();
   } catch (e: any) {
-    expect(e.message).toEqual(
-      "The discriminator value could not be extracted from all the provided schemas"
-    );
+    expect(e.message).toContain("could not be extracted");
   }
 });
 
@@ -210,9 +141,7 @@ test("wrong schema - duplicate discriminator values", () => {
     ]);
     throw new Error();
   } catch (e: any) {
-    expect(e.message).toEqual(
-      "Some of the discriminator values are not unique"
-    );
+    expect(e.message).toContain("has duplicate value");
   }
 });
 

deno/lib/types.ts

@@ -18,7 +18,7 @@ import {
   ParsePath,
   ParseReturnType,
   ParseStatus,
-  SyncParseReturnType
+  SyncParseReturnType,
 } from "./helpers/parseUtil.ts";
 import { partialUtil } from "./helpers/partialUtil.ts";
 import { Primitive } from "./helpers/typeAliases.ts";
@@ -30,7 +30,7 @@ import {
   ZodError,
   ZodErrorMap,
   ZodIssue,
-  ZodIssueCode
+  ZodIssueCode,
 } from "./ZodError.ts";
 
 ///////////////////////////////////////
@@ -2137,80 +2137,46 @@ export class ZodUnion<T extends ZodUnionOptions> extends ZodType<
 /////////////////////////////////////////////////////
 /////////////////////////////////////////////////////
 
-type ZodSourceType<T extends ZodTypeAny> = T extends ZodLazy<infer U>
-  ? ZodSourceType<U>
-  : T extends ZodEffects<infer U>
-  ? ZodSourceType<U>
-  : T;
-
-type Prev = [never, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
-
-type ZodOriginType<T extends ZodTypeAny, D extends Prev[number] = 9> = [
-  D
-] extends [never]
-  ? never
-  :
-      | T
-      | ZodEffects<ZodOriginType<T, Prev[D]>>
-      | ZodLazy<ZodOriginType<T, Prev[D]>>;
-
-function getSourceType<T extends ZodTypeAny>(type: T): ZodSourceType<T> {
-  if (type._def.typeName === ZodFirstPartyTypeKind.ZodLazy) {
-    return getSourceType(
-      (type as unknown as ZodLazy<ZodTypeAny>).schema
-    ) as ZodSourceType<T>;
-  } else if (type._def.typeName === ZodFirstPartyTypeKind.ZodEffects) {
-    return getSourceType(
-      (type as unknown as ZodEffects<ZodTypeAny>).sourceType()
-    ) as ZodSourceType<T>;
+const getDiscriminator = <T extends ZodTypeAny>(
+  type: T
+): Primitive[] | null => {
+  if (type instanceof ZodLazy) {
+    return getDiscriminator(type.schema);
+  } else if (type instanceof ZodEffects) {
+    return getDiscriminator(type.innerType());
+  } else if (type instanceof ZodLiteral) {
+    return [type.value];
+  } else if (type instanceof ZodEnum) {
+    return type.options;
+  } else if (type instanceof ZodUndefined) {
+    return [undefined];
+  } else if (type instanceof ZodNull) {
+    return [null];
   } else {
-    return type as ZodSourceType<T>;
+    return null;
   }
-}
-
-type ZodDiscriminatedUnionOptionBase<
-  Discriminator extends string,
-  DiscriminatorValue extends Primitive
-> = ZodObject<
-  { [key in Discriminator]: ZodLiteral<DiscriminatorValue> } & ZodRawShape,
-  any,
-  any
->;
-
-type ZodDiscriminatedUnionType<Discriminator extends string> = Record<
-  string,
-  any
-> & {
-  [key in Discriminator]: Primitive;
 };
 
-export type ZodDiscriminatedUnionOption<
-  Discriminator extends string,
-  DiscriminatorValue extends Primitive
-> =
-  | ZodOriginType<
-      ZodDiscriminatedUnionOptionBase<Discriminator, DiscriminatorValue>
-    >
-  | ZodType<any, any, ZodDiscriminatedUnionType<Discriminator>>;
+export type ZodDiscriminatedUnionOption<Discriminator extends string> =
+  ZodObject<{ [key in Discriminator]: ZodTypeAny } & ZodRawShape, any, any>;
 
 export interface ZodDiscriminatedUnionDef<
   Discriminator extends string,
-  DiscriminatorValue extends Primitive,
-  Option extends ZodDiscriminatedUnionOption<Discriminator, DiscriminatorValue>
+  Options extends ZodDiscriminatedUnionOption<any>[]
 > extends ZodTypeDef {
   discriminator: Discriminator;
-  options: Map<DiscriminatorValue, Option>;
+  options: Options;
+  optionsMap: Map<Primitive, ZodDiscriminatedUnionOption<any>>;
   typeName: ZodFirstPartyTypeKind.ZodDiscriminatedUnion;
 }
 
 export class ZodDiscriminatedUnion<
   Discriminator extends string,
-  DiscriminatorValue extends Primitive,
-  Option extends ZodDiscriminatedUnionOption<Discriminator, DiscriminatorValue>
+  Options extends ZodDiscriminatedUnionOption<Discriminator>[]
 > extends ZodType<
-  output<Option>,
-  ZodDiscriminatedUnionDef<Discriminator, DiscriminatorValue, Option>,
-  input<Option>
+  output<Options[number]>,
+  ZodDiscriminatedUnionDef<Discriminator, Options>,
+  input<Options[number]>
 > {
   _parse(input: ParseInput): ParseReturnType<this["_output"]> {
     const { ctx } = this._processInputParams(input);
@@ -2225,13 +2191,13 @@ export class ZodDiscriminatedUnion<
     }
 
     const discriminator = this.discriminator;
-    const discriminatorValue: DiscriminatorValue = ctx.data[discriminator];
-    const option = this.options.get(discriminatorValue);
+    const discriminatorValue: string = ctx.data[discriminator];
+    const option = this.optionsMap.get(discriminatorValue);
 
     if (!option) {
       addIssueToContext(ctx, {
         code: ZodIssueCode.invalid_union_discriminator,
-        options: this.validDiscriminatorValues,
+        options: Array.from(this.optionsMap.keys()),
         path: [discriminator],
       });
       return INVALID;
@@ -2242,28 +2208,28 @@ export class ZodDiscriminatedUnion<
         data: ctx.data,
         path: ctx.path,
         parent: ctx,
-      });
+      }) as any;
     } else {
       return option._parseSync({
         data: ctx.data,
         path: ctx.path,
         parent: ctx,
-      });
+      }) as any;
     }
   }
 
   get discriminator() {
     return this._def.discriminator;
   }
 
-  get validDiscriminatorValues() {
-    return Array.from(this.options.keys());
-  }
-
   get options() {
     return this._def.options;
   }
 
+  get optionsMap() {
+    return this._def.optionsMap;
+  }
+
   /**
    * The constructor of the discriminated union schema. Its behaviour is very similar to that of the normal z.union() constructor.
    * However, it only allows a union of objects, all of which need to share a discriminator property. This property must
@@ -2274,48 +2240,45 @@ export class ZodDiscriminatedUnion<
    */
   static create<
     Discriminator extends string,
-    DiscriminatorValue extends Primitive,
     Types extends [
-      ZodDiscriminatedUnionOption<Discriminator, DiscriminatorValue>,
-      ZodDiscriminatedUnionOption<Discriminator, DiscriminatorValue>,
-      ...ZodDiscriminatedUnionOption<Discriminator, DiscriminatorValue>[]
+      ZodDiscriminatedUnionOption<Discriminator>,
+      ...ZodDiscriminatedUnionOption<Discriminator>[]
     ]
   >(
     discriminator: Discriminator,
-    types: Types,
+    options: Types,
     params?: RawCreateParams
-  ): ZodDiscriminatedUnion<Discriminator, DiscriminatorValue, Types[number]> {
+  ): ZodDiscriminatedUnion<Discriminator, Types> {
     // Get all the valid discriminator values
-    const options: Map<DiscriminatorValue, Types[number]> = new Map();
-
-    try {
-      types.forEach((type) => {
-        const discriminatorValue = getSourceType(
-          type as ZodOriginType<
-            ZodDiscriminatedUnionOptionBase<Discriminator, DiscriminatorValue>
-          >
-        ).shape[discriminator].value;
-        options.set(discriminatorValue, type);
-      });
-    } catch (e) {
-      throw new Error(
-        "The discriminator value could not be extracted from all the provided schemas"
-      );
-    }
-
-    // Assert that all the discriminator values are unique
-    if (options.size !== types.length) {
-      throw new Error("Some of the discriminator values are not unique");
+    const optionsMap: Map<Primitive, Types[number]> = new Map();
+
+    // try {
+    for (const type of options) {
+      const discriminatorValues = getDiscriminator(type.shape[discriminator]);
+      if (!discriminatorValues) {
+        throw new Error(
+          `A discriminator value for key \`${discriminator}\`could not be extracted from all schema options`
+        );
+      }
+      for (const value of discriminatorValues) {
+        if (optionsMap.has(value)) {
+          throw new Error(
+            `Discriminator property ${discriminator} has duplicate value ${value}`
+          );
+        }
+        optionsMap.set(value, type);
+      }
     }
 
     return new ZodDiscriminatedUnion<
       Discriminator,
-      DiscriminatorValue,
-      Types[number]
+      // DiscriminatorValue,
+      Types
     >({
       typeName: ZodFirstPartyTypeKind.ZodDiscriminatedUnion,
       discriminator,
       options,
+      optionsMap,
       ...processCreateParams(params),
     });
   }
@@ -3911,7 +3874,7 @@ export type ZodFirstPartySchemaTypes =
   | ZodArray<any, any>
   | ZodObject<any, any, any, any, any>
   | ZodUnion<any>
-  | ZodDiscriminatedUnion<any, any, any>
+  | ZodDiscriminatedUnion<any, any>
   | ZodIntersection<any, any>
   | ZodTuple<any, any>
   | ZodRecord<any, any>