.extend on a ZodUnion of objects

[rijenkii]

Recently I really needed a method for .extending all elements of a union and returning a new union.
Because I am a newb with TypeScript, this has taken me multiple days, but I was able to concoct something that I think works:

import * as z from "zod";

function zodUnionExtend<
  Us extends [z.AnyZodObject, z.AnyZodObject, ...z.AnyZodObject[]],
  S extends z.ZodRawShape
>(
  union: z.ZodUnion<Us>,
  shape: S
): z.ZodUnion<{
  [K in keyof Us]: z.ZodObject<
    z.objectUtil.extendShape<Us[K]["shape"], S>,
    Us[K]["_def"]["unknownKeys"],
    Us[K]["_def"]["catchall"]
  >;
}> {
  /* eslint-disable-next-line @typescript-eslint/ban-ts-comment */
  // @ts-ignore
  return z.union(union.options.map((o) => o.extend(shape)));
}

TypeScript Playground

I tried getting return type of .extend through ReturnType, but in the end I did not succeed and just copied the signature from the source code.

[mmvsk]

Use ZodIntersection:

const CircleSchema = Z.object({
    radius: Z.number(),
});

const RectangleSchema = Z.object({
    sideA: Z.number(),
    sideB: Z.number(),
});

const ShapeSchema = Z.discriminatedUnion("type", [
    CircleSchema.extend({ type: Z.literal("circle") }),
    RectangleSchema.extend({ type: Z.literal("rectangle") }),
]);

const PositionSchema = Z.object({ x: Z.number(), y: Z.number() });

const PlanObjectSchema = Z.intersection(ShapeSchema, PositionSchema);

type Shape = Z.output<typeof ShapeSchema>;
type Position = Z.output<typeof PositionSchema>;
type PlanObject = Z.output<typeof PlanObjectSchema>;

function CreatePlanObject(shape: Shape, position: Position): PlanObject {
    return {
        ...shape,
        ...position,
    };
}

[andho]

Sadly this will lose function of ZodObject like pick and omit according to docs, but seems like the only way for now.