#1171 support for refine, superRefine, transform and lazy in discriminatedUnion (#1290)

* #1171

* fix tests

* add superRefine in tests

* add support for lazy

* fix typings

* fixe typings for asserted lazy

* fix

* clean console.log from debug

* Clean up discriminatedUnion

* Fix deno test

Co-authored-by: Colin McDonnell <[email protected]>

Author: roblabat

deno/lib/__tests__/discriminatedUnions.test.ts

@@ -25,6 +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) }),
+    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 });
@@ -126,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.includes("could not be extracted")).toBe(true);
   }
 });
 
@@ -140,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.includes("has duplicate value")).toEqual(true);
   }
 });
 

deno/lib/types.ts

@@ -2297,33 +2297,46 @@ export class ZodUnion<T extends ZodUnionOptions> extends ZodType<
 /////////////////////////////////////////////////////
 /////////////////////////////////////////////////////
 
-export type ZodDiscriminatedUnionOption<
-  Discriminator extends string,
-  DiscriminatorValue extends Primitive
-> = ZodObject<
-  { [key in Discriminator]: ZodLiteral<DiscriminatorValue> } & ZodRawShape,
-  any,
-  any
->;
+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 null;
+  }
+};
+
+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<
-  Option["_output"],
-  ZodDiscriminatedUnionDef<Discriminator, DiscriminatorValue, Option>,
-  Option["_input"]
+  output<Options[number]>,
+  ZodDiscriminatedUnionDef<Discriminator, Options>,
+  input<Options[number]>
 > {
   _parse(input: ParseInput): ParseReturnType<this["_output"]> {
     const { ctx } = this._processInputParams(input);
@@ -2338,13 +2351,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;
@@ -2355,28 +2368,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
@@ -2387,44 +2400,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 = type.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),
     });
   }
@@ -3570,13 +3584,19 @@ export interface ZodEffectsDef<T extends ZodTypeAny = ZodTypeAny>
 
 export class ZodEffects<
   T extends ZodTypeAny,
-  Output = T["_output"],
-  Input = T["_input"]
+  Output = output<T>,
+  Input = input<T>
 > extends ZodType<Output, ZodEffectsDef<T>, Input> {
   innerType() {
     return this._def.schema;
   }
 
+  sourceType(): T {
+    return this._def.schema._def.typeName === ZodFirstPartyTypeKind.ZodEffects
+      ? (this._def.schema as unknown as ZodEffects<T>).sourceType()
+      : (this._def.schema as T);
+  }
+
   _parse(input: ParseInput): ParseReturnType<this["_output"]> {
     const { status, ctx } = this._processInputParams(input);
 
@@ -4161,7 +4181,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>

src/__tests__/discriminatedUnions.test.ts

@@ -24,6 +24,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) }),
+    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 });
@@ -125,9 +128,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.includes("could not be extracted")).toBe(true);
   }
 });
 
@@ -139,9 +140,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.includes("has duplicate value")).toEqual(true);
   }
 });