feat(enum): return enum from object keys (#1216)

* feat(enum): return enum from object keys

* docs(enum): add enum documentation

* docs: fix documentation

* Change enum to keyof

* Fix lint

* Run prettier

* Fix Deno tests

Co-authored-by: Colin McDonnell <colinmcd@alum.mit.edu>

README.md

@@ -63,6 +63,7 @@
 - [Nullables](#nullables)
 - [Objects](#objects)
   - [.shape](#shape)
+  - [.enum](#enum)
   - [.extend](#extend)
   - [.merge](#merge)
   - [.pick/.omit](#pickomit)
@@ -782,6 +783,15 @@ Dog.shape.name; // => string schema
 Dog.shape.age; // => number schema
 ```
 
+### `.keyof`
+
+Use `.key` to create a `ZodEnum` schema from the keys of an object schema.
+
+```ts
+const keySchema = Dog.keyof();
+keySchema; // ZodEnum<["name", "age"]>
+```
+
 ### `.extend`
 
 You can add additional fields to an object schema with the `.extend` method.

deno/lib/README.md

@@ -63,6 +63,7 @@
 - [Nullables](#nullables)
 - [Objects](#objects)
   - [.shape](#shape)
+  - [.enum](#enum)
   - [.extend](#extend)
   - [.merge](#merge)
   - [.pick/.omit](#pickomit)
@@ -782,6 +783,15 @@ Dog.shape.name; // => string schema
 Dog.shape.age; // => number schema
 ```
 
+### `.keyof`
+
+Use `.key` to create a `ZodEnum` schema from the keys of an object schema.
+
+```ts
+const keySchema = Dog.keyof();
+keySchema; // ZodEnum<["name", "age"]>
+```
+
 ### `.extend`
 
 You can add additional fields to an object schema with the `.extend` method.

deno/lib/__tests__/object.test.ts

@@ -238,6 +238,23 @@ test("inferred unioned object type with optional properties", async () => {
   f1;
 });
 
+test("inferred enum type", async () => {
+  const Enum = z.object({ a: z.string(), b: z.string().optional() }).keyof();
+
+  expect(Enum.Values).toEqual({
+    a: "a",
+    b: "b",
+  });
+  expect(Enum.enum).toEqual({
+    a: "a",
+    b: "b",
+  });
+  expect(Enum._def.values).toEqual(["a", "b"]);
+  type Enum = z.infer<typeof Enum>;
+  const f1: util.AssertEqual<Enum, "a" | "b"> = true;
+  f1;
+});
+
 test("inferred partial object type with optional properties", async () => {
   const Partial = z
     .object({ a: z.string(), b: z.string().optional() })

deno/lib/__tests__/recursive.test.ts

@@ -1,6 +1,7 @@
 // @ts-ignore TS6133
 import { expect } from "https://deno.land/x/expect@v0.2.6/mod.ts";
 const test = Deno.test;
+
 import { z } from "../index.ts";
 
 interface Category {

deno/lib/helpers/enumUtil.ts

@@ -0,0 +1,19 @@
+export namespace enumUtil {
+  type UnionToIntersectionFn<T> = (
+    T extends unknown ? (k: () => T) => void : never
+  ) extends (k: infer Intersection) => void
+    ? Intersection
+    : never;
+
+  type GetUnionLast<T> = UnionToIntersectionFn<T> extends () => infer Last
+    ? Last
+    : never;
+
+  type UnionToTuple<T, Tuple extends unknown[] = []> = [T] extends [never]
+    ? Tuple
+    : UnionToTuple<Exclude<T, GetUnionLast<T>>, [GetUnionLast<T>, ...Tuple]>;
+
+  type CastToStringTuple<T> = T extends [string, ...string[]] ? T : never;
+
+  export type UnionToTupleString<T> = CastToStringTuple<UnionToTuple<T>>;
+}

deno/lib/types.ts

@@ -1,3 +1,4 @@
+import { enumUtil } from "./helpers/enumUtil.ts";
 import { errorUtil } from "./helpers/errorUtil.ts";
 import {
   addIssueToContext,
@@ -1932,6 +1933,12 @@ export class ZodObject<
     }) as any;
   }
 
+  keyof(): ZodEnum<enumUtil.UnionToTupleString<keyof T>> {
+    return createZodEnum(
+      util.objectKeys(this.shape) as [string, ...string[]]
+    ) as any;
+  }
+
   static create = <T extends ZodRawShape>(
     shape: T,
     params?: RawCreateParams

src/__tests__/object.test.ts

@@ -237,6 +237,23 @@ test("inferred unioned object type with optional properties", async () => {
   f1;
 });
 
+test("inferred enum type", async () => {
+  const Enum = z.object({ a: z.string(), b: z.string().optional() }).keyof();
+
+  expect(Enum.Values).toEqual({
+    a: "a",
+    b: "b",
+  });
+  expect(Enum.enum).toEqual({
+    a: "a",
+    b: "b",
+  });
+  expect(Enum._def.values).toEqual(["a", "b"]);
+  type Enum = z.infer<typeof Enum>;
+  const f1: util.AssertEqual<Enum, "a" | "b"> = true;
+  f1;
+});
+
 test("inferred partial object type with optional properties", async () => {
   const Partial = z
     .object({ a: z.string(), b: z.string().optional() })

src/__tests__/recursive.test.ts

@@ -1,5 +1,6 @@
 // @ts-ignore TS6133
 import { test } from "@jest/globals";
+
 import { z } from "..";
 
 interface Category {

src/helpers/enumUtil.ts

@@ -0,0 +1,19 @@
+export namespace enumUtil {
+  type UnionToIntersectionFn<T> = (
+    T extends unknown ? (k: () => T) => void : never
+  ) extends (k: infer Intersection) => void
+    ? Intersection
+    : never;
+
+  type GetUnionLast<T> = UnionToIntersectionFn<T> extends () => infer Last
+    ? Last
+    : never;
+
+  type UnionToTuple<T, Tuple extends unknown[] = []> = [T] extends [never]
+    ? Tuple
+    : UnionToTuple<Exclude<T, GetUnionLast<T>>, [GetUnionLast<T>, ...Tuple]>;
+
+  type CastToStringTuple<T> = T extends [string, ...string[]] ? T : never;
+
+  export type UnionToTupleString<T> = CastToStringTuple<UnionToTuple<T>>;
+}

src/types.ts

@@ -1,3 +1,4 @@
+import { enumUtil } from "./helpers/enumUtil";
 import { errorUtil } from "./helpers/errorUtil";
 import {
   addIssueToContext,
@@ -1932,6 +1933,12 @@ export class ZodObject<
     }) as any;
   }
 
+  keyof(): ZodEnum<enumUtil.UnionToTupleString<keyof T>> {
+    return createZodEnum(
+      util.objectKeys(this.shape) as [string, ...string[]]
+    ) as any;
+  }
+
   static create = <T extends ZodRawShape>(
     shape: T,
     params?: RawCreateParams