feat(v4): add z.getDiscriminatedOption (#5947)

Adds `z.getDiscriminatedOption(union, value)`, which returns the option of a
discriminated union whose discriminator claims `value`. The result is narrowed
to the single matching option, including when one option declares several
values via `z.literal([...])`.

Defined in core and re-exported from both classic and mini, so both get the
identical function. The derived map is cached in the existing `_zod.bag`, so no
new internals property is added and nothing is computed until the util is first
called. The parse path is untouched.

Closes the by-value lookup half of #5086.

Co-authored-by: Alessandro Colace <ale.colace@gmail.com>

Author: dokson

packages/zod/src/v4/classic/external.ts

@@ -27,6 +27,7 @@ export {
   compile,
   ZodCompileAsyncError,
   ZodCompileUnsupportedError,
+  getDiscriminatedOption,
 } from "../core/index.js";
 export { toJSONSchema } from "../core/json-schema-processors.js";
 export { fromJSONSchema } from "./from-json-schema.js";

packages/zod/src/v4/classic/tests/discriminated-unions.test.ts

@@ -740,6 +740,57 @@ test("encode with codec discriminator", () => {
   expect(encoded).toEqual({ type: 1, value: "hello" });
 });
 
+test("getDiscriminatedOption", () => {
+  const fruit = z.object({ type: z.literal("fruit"), seeds: z.boolean() });
+  const veg = z.object({ type: z.literal("vegetable"), leafy: z.boolean() });
+  const schema = z.discriminatedUnion("type", [fruit, veg]);
+
+  expect(z.getDiscriminatedOption(schema, "fruit")).toBe(fruit);
+  expect(z.getDiscriminatedOption(schema, "vegetable")).toBe(veg);
+
+  // The result is narrowed to the one member, not the union of all of them.
+  expectTypeOf(z.getDiscriminatedOption(schema, "fruit")).toEqualTypeOf<typeof fruit>();
+  expectTypeOf(z.getDiscriminatedOption(schema, "vegetable")).toEqualTypeOf<typeof veg>();
+  // @ts-expect-error — "unknown" is not a declared discriminator value
+  z.getDiscriminatedOption(schema, "unknown");
+});
+
+test("getDiscriminatedOption — multi-value members and non-string discriminators", () => {
+  const a = z.object({ type: z.literal(["x", "y"]), payload: z.string() });
+  const b = z.object({ type: z.literal("z"), payload: z.number() });
+  const multi = z.discriminatedUnion("type", [a, b]);
+  expect(z.getDiscriminatedOption(multi, "x")).toBe(a);
+  expect(z.getDiscriminatedOption(multi, "y")).toBe(a);
+  expect(z.getDiscriminatedOption(multi, "z")).toBe(b);
+
+  const num = z.object({ type: z.literal(1) });
+  const bool = z.object({ type: z.literal(true) });
+  const nul = z.object({ type: z.null() });
+  const mixed = z.discriminatedUnion("type", [num, bool, nul]);
+  expect(z.getDiscriminatedOption(mixed, 1)).toBe(num);
+  expect(z.getDiscriminatedOption(mixed, true)).toBe(bool);
+  expect(z.getDiscriminatedOption(mixed, null)).toBe(nul);
+
+  // An omittable discriminator claims undefined, so it resolves like any other value.
+  const opt = z.object({ type: z.literal("a").optional(), a: z.string() });
+  const req = z.object({ type: z.literal("b"), b: z.string() });
+  expect(z.getDiscriminatedOption(z.discriminatedUnion("type", [opt, req]), undefined)).toBe(opt);
+});
+
+test("getDiscriminatedOption caches in the bag and costs nothing until called", () => {
+  const schema = z.discriminatedUnion("type", [z.object({ type: z.literal("a") }), z.object({ type: z.literal("b") })]);
+
+  expect(schema._zod.bag.optionsMap).toBeUndefined();
+  expect(z.getDiscriminatedOption(schema, "a")).toBe(schema.options[0]);
+  const map = schema._zod.bag.optionsMap;
+  expect(map).toBeInstanceOf(Map);
+  expect(z.getDiscriminatedOption(schema, "b")).toBe(schema.options[1]);
+  expect(schema._zod.bag.optionsMap).toBe(map);
+
+  // A clone recomputes rather than inheriting the cache.
+  expect(schema.clone()._zod.bag.optionsMap).toBeUndefined();
+});
+
 test.each(["__proto__", "constructor", "toString", "hasOwnProperty", "valueOf"])(
   "Object.prototype discriminator name: %s",
   (key) => {

packages/zod/src/v4/core/schemas.ts

@@ -2482,6 +2482,49 @@ export interface $ZodDiscriminatedUnionInternals<
 > extends $ZodUnionInternals<Options> {
   def: $ZodDiscriminatedUnionDef<Options, Disc>;
   propValues: util.PropValues;
+  bag: util.LoosePartial<{
+    optionsMap: Map<util.Primitive, $ZodType>;
+  }>;
+}
+
+/** The discriminator values a member of `Options` can declare. An omittable discriminator contributes `undefined`, matching what `propValues` claims for it. */
+export type $DiscriminatorValue<Options extends readonly SomeType[], Disc extends string> = {
+  [I in keyof Options]: Options[I] extends { _zod: { output: infer Out } }
+    ? Disc extends keyof Out
+      ? Out[Disc]
+      : never
+    : never;
+}[number];
+
+/** The member of `Options` that declares `V`. */
+export type $DiscriminatedOption<Options extends readonly SomeType[], Disc extends string, V> = {
+  [I in keyof Options]: Options[I] extends { _zod: { output: infer Out } }
+    ? Disc extends keyof Out
+      ? V extends Out[Disc]
+        ? Options[I]
+        : never
+      : never
+    : never;
+}[number];
+
+/** Returns the option of `union` whose discriminator claims `value`. */
+export function getDiscriminatedOption<
+  Options extends readonly SomeType[],
+  Disc extends string,
+  const V extends $DiscriminatorValue<Options, Disc>,
+>(union: $ZodDiscriminatedUnion<Options, Disc>, value: V): $DiscriminatedOption<Options, Disc, V> {
+  const internals = union._zod;
+  let map = internals.bag.optionsMap;
+  if (!map) {
+    map = new Map();
+    const { options, discriminator } = internals.def;
+    for (const option of options as unknown as readonly $ZodType[]) {
+      // First declaration wins, matching the order the parse path resolves a duplicate in.
+      for (const v of option._zod.propValues?.[discriminator] ?? []) if (!map.has(v)) map.set(v, option);
+    }
+    internals.bag.optionsMap = map;
+  }
+  return map.get(value as util.Primitive) as any;
 }
 
 export interface $ZodDiscriminatedUnion<

packages/zod/src/v4/mini/external.ts

@@ -26,6 +26,7 @@ export {
   compile,
   ZodCompileAsyncError,
   ZodCompileUnsupportedError,
+  getDiscriminatedOption,
 } from "../core/index.js";
 export { toJSONSchema } from "../core/json-schema-processors.js";
 

packages/zod/src/v4/mini/tests/index.test.ts

@@ -1032,3 +1032,13 @@ test("type narrowing works with type property", () => {
     expect(arraySchema.def.element).toBeDefined();
   }
 });
+
+test("getDiscriminatedOption", () => {
+  const a = z.object({ type: z.literal("a"), x: z.string() });
+  const b = z.object({ type: z.literal("b"), y: z.number() });
+  const schema = z.discriminatedUnion("type", [a, b]);
+
+  expect(z.getDiscriminatedOption(schema, "a")).toBe(a);
+  expect(z.getDiscriminatedOption(schema, "b")).toBe(b);
+  expectTypeOf(z.getDiscriminatedOption(schema, "a")).toEqualTypeOf<typeof a>();
+});