feat: Add extract/exclude methods to ZodEnum (#1652)

* Add extract/exclude to ZodEnum * Add specs for extract/exclude * Fix types that were failing build * Add test to extract exclude --------- Co-authored-by: santosmarco <eu@marco.rio.br>
colinhacks · Feb 8, 2023 · 37d83e8 · 37d83e8
1 parent 1e74925
commit 37d83e8
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 0 deletions.
diff --git a/deno/lib/__tests__/enum.test.ts b/deno/lib/__tests__/enum.test.ts
@@ -41,3 +41,20 @@ test("error params", () => {
     expect(result.error.issues[0].message).toEqual("REQUIRED");
   }
 });
+
+test("extract/exclude", () => {
+  const foods = ["Pasta", "Pizza", "Tacos", "Burgers", "Salad"] as const;
+  const FoodEnum = z.enum(foods);
+  const ItalianEnum = FoodEnum.extract(["Pasta", "Pizza"]);
+  const UnhealthyEnum = FoodEnum.exclude(["Salad"]);
+  const EmptyFoodEnum = FoodEnum.exclude(foods);
+
+  util.assertEqual<z.infer<typeof ItalianEnum>, "Pasta" | "Pizza">(true);
+  util.assertEqual<
+    z.infer<typeof UnhealthyEnum>,
+    "Pasta" | "Pizza" | "Tacos" | "Burgers"
+  >(true);
+  // @ts-expect-error TS2344
+  util.assertEqual<typeof EmptyFoodEnum, z.ZodEnum<[]>>(true);
+  util.assertEqual<z.infer<typeof EmptyFoodEnum>, never>(true);
+});
diff --git a/deno/lib/types.ts b/deno/lib/types.ts
@@ -3547,6 +3547,14 @@ export interface ZodEnumDef<T extends EnumValues = EnumValues>
 
 export type Writeable<T> = { -readonly [P in keyof T]: T[P] };
 
+export type FilterEnum<Values, ToExclude> = Values extends []
+  ? []
+  : Values extends [infer Head, ...infer Rest]
+  ? Head extends ToExclude
+    ? FilterEnum<Rest, ToExclude>
+    : [Head, ...FilterEnum<Rest, ToExclude>]
+  : never;
+
 function createZodEnum<U extends string, T extends Readonly<[U, ...U[]]>>(
   values: T,
   params?: RawCreateParams
@@ -3621,6 +3629,23 @@ export class ZodEnum<T extends [string, ...string[]]> extends ZodType<
     return enumValues as any;
   }
 
+  extract<ToExtract extends readonly [T[number], ...T[number][]]>(
+    values: ToExtract
+  ) {
+    return ZodEnum.create(values);
+  }
+
+  exclude<ToExclude extends readonly [T[number], ...T[number][]]>(
+    values: ToExclude
+  ) {
+    return ZodEnum.create(
+      this.options.filter((opt) => !values.includes(opt)) as FilterEnum<
+        T,
+        ToExclude[number]
+      >
+    );
+  }
+
   static create = createZodEnum;
 }
 

diff --git a/src/__tests__/enum.test.ts b/src/__tests__/enum.test.ts
@@ -40,3 +40,20 @@ test("error params", () => {
     expect(result.error.issues[0].message).toEqual("REQUIRED");
   }
 });
+
+test("extract/exclude", () => {
+  const foods = ["Pasta", "Pizza", "Tacos", "Burgers", "Salad"] as const;
+  const FoodEnum = z.enum(foods);
+  const ItalianEnum = FoodEnum.extract(["Pasta", "Pizza"]);
+  const UnhealthyEnum = FoodEnum.exclude(["Salad"]);
+  const EmptyFoodEnum = FoodEnum.exclude(foods);
+
+  util.assertEqual<z.infer<typeof ItalianEnum>, "Pasta" | "Pizza">(true);
+  util.assertEqual<
+    z.infer<typeof UnhealthyEnum>,
+    "Pasta" | "Pizza" | "Tacos" | "Burgers"
+  >(true);
+  // @ts-expect-error TS2344
+  util.assertEqual<typeof EmptyFoodEnum, z.ZodEnum<[]>>(true);
+  util.assertEqual<z.infer<typeof EmptyFoodEnum>, never>(true);
+});
diff --git a/src/types.ts b/src/types.ts
@@ -3547,6 +3547,14 @@ export interface ZodEnumDef<T extends EnumValues = EnumValues>
 
 export type Writeable<T> = { -readonly [P in keyof T]: T[P] };
 
+export type FilterEnum<Values, ToExclude> = Values extends []
+  ? []
+  : Values extends [infer Head, ...infer Rest]
+  ? Head extends ToExclude
+    ? FilterEnum<Rest, ToExclude>
+    : [Head, ...FilterEnum<Rest, ToExclude>]
+  : never;
+
 function createZodEnum<U extends string, T extends Readonly<[U, ...U[]]>>(
   values: T,
   params?: RawCreateParams
@@ -3621,6 +3629,23 @@ export class ZodEnum<T extends [string, ...string[]]> extends ZodType<
     return enumValues as any;
   }
 
+  extract<ToExtract extends readonly [T[number], ...T[number][]]>(
+    values: ToExtract
+  ) {
+    return ZodEnum.create(values);
+  }
+
+  exclude<ToExclude extends readonly [T[number], ...T[number][]]>(
+    values: ToExclude
+  ) {
+    return ZodEnum.create(
+      this.options.filter((opt) => !values.includes(opt)) as FilterEnum<
+        T,
+        ToExclude[number]
+      >
+    );
+  }
+
   static create = createZodEnum;
 }