Zod Deep Pick

zod_deep_pick.ts

function isZodObject(schema: z.ZodTypeAny): schema is z.AnyZodObject {
  if (schema._def.typeName === "ZodObject") return true;

  return false;
}

function isZodArray(schema: z.ZodTypeAny): schema is z.ZodArray<any> {
  if (schema._def.typeName === "ZodArray") return true;

  return false;
}

function pickObject(schema: z.ZodTypeAny, path: string): z.ZodTypeAny {
  if (!isZodObject(schema)) throw Error("Not a zod object");

  const newSchema = schema.shape?.[path];
  if (!newSchema) throw Error(`${path} does not exist on schema with keys: ${Object.keys(schema.shape)}`);

  return newSchema;
}

function pickArray(schema: z.ZodTypeAny): z.ZodTypeAny {
  if (!isZodArray(schema)) throw Error("Not a Zod Array");

  const newSchema = schema?.element;
  if (!newSchema) throw Error("No element on Zod Array");

  return newSchema;
}

function zodDeepPick(schema: z.ZodTypeAny, propertyPath: string): z.ZodTypeAny {
  if (propertyPath === "") return schema;

  const numberRegex = new RegExp("[[0-9]+]");

  const arrayIndex = propertyPath.search(numberRegex);
  const objectIndex = propertyPath.indexOf(".");

  const matchedArray = arrayIndex !== -1;
  const matchedObject = objectIndex !== -1;

  if ((matchedArray && matchedObject && arrayIndex < objectIndex) || (matchedArray && !matchedObject)) {
    const arraySplit = propertyPath.split(numberRegex);
    const restArray = arraySplit.slice(1, arraySplit.length).join("[0]");

    if (arrayIndex !== 0) {
      return zodDeepPick(pickObject(schema, arraySplit[0]), `[0]${restArray}`);
    }

    return zodDeepPick(
      pickArray(schema),
      restArray.charAt(0) === "." ? restArray.slice(1, restArray.length) : restArray
    );
  }

  if (matchedObject) {
    const objectSplit = propertyPath.split(".");
    const restObject = objectSplit.slice(1, objectSplit.length).join(".");

    return zodDeepPick(pickObject(schema, objectSplit[0]), restObject);
  }

  return pickObject(schema, propertyPath);
}
                                                                 
// USAGE
                                                                 
const schema = z.object({
  fake: z.object({
    thing: z.array(z.string()),
    foo: z.array(z.array(z.string().email())),
    peanut: z.array(
      z.object({
        butter: z.number(),
        test: z.array(z.boolean()),
        testing: z.array(
          z.object({
            hi: z.array(z.string()),
          })
        ),
      })
    ),
  }),
});

zodDeepPick(schema, "fake.thing[69]") // z.string()
zodDeepPick(schema, "fake.foo[69][23]") // z.string().email()
zodDeepPick(schema, "fake.peanut[32].butter") // z.number()
zodDeepPick(schema, "fake.peanut[0].test[420]") // z.boolean()
zodDeepPick(schema, "fake.peanut[32].testing[432].hi") // z.array(z.string())
zodDeepPick(schema, "fake.notreal") // throws Error: notreal does not exist on schema with keys: thing,foo,peanut
zodDeepPick(schema, "fake[0]") // throws Error: Not a Zod Array