Extend zod functionality by wrapping it

zpp.ts

/* eslint-disable @typescript-eslint/ban-ts-comment */
import endent from "endent";
import { type z, type ZodRawShape } from "zod";

export const zpp = <Out, Def extends z.ZodTypeDef, In>(
  schema: z.Schema<Out, Def, In>
) => {
  const newFunc = (obj: z.infer<typeof schema>) => {
    return obj;
  };

  const parseJson = (data: unknown) => {
    if (typeof data === "string") {
      return schema.parse(JSON.parse(data));
    }
    if (typeof data === "object") {
      return schema.parse(data);
    }

    throw new Error("Invalid data type");
  };

  const stringify = (data: z.infer<typeof schema>) => {
    return JSON.stringify(data);
  };

  const parseJsonSafe = (data: unknown) => {
    if (typeof data === "string") {
      return schema.safeParse(JSON.parse(data));
    }
    if (typeof data === "object") {
      return schema.safeParse(data);
    }

    throw new Error("Invalid data type");
  };

  const outputPrompt = (objectSchema: z.ZodObject<ZodRawShape>) => {
    let output = "{";

    const keys = Object.keys(objectSchema.shape);

    for (const key of keys) {
      const shape = objectSchema.shape[key];
      if (!shape) throw new Error("Invalid Shape");

      const def = shape._def;

      if (def.typeName === "ZodObject") {
        // @ts-expect-error
        output = `${output}\n"${key}": ${outputPrompt(shape)}`;
        continue;
      }

      let typeName = def.typeName;

      if (typeName === "ZodOptional") {
        typeName = def.innerType._def.typeName;
      }

      const describe: string | undefined = def.description;

      switch (typeName) {
        case "ZodString":
          output = `${output}\n"${key}": ${
            describe ? `"${describe}"` : "string"
          },`;
          break;
        case "ZodNumber":
          output = `${output}\n"${key}": number${
            describe ? ` // ${describe}` : ""
          },`;
          break;
        case "ZodBoolean":
          output = `${output}\n"${key}": boolean${
            describe ? ` // ${describe}` : ""
          },`;
          break;
      }
    }

    // Replace trailing comma
    output = output.replace(/,\s*$/, "");

    output = endent`${output}\n}`;

    return output;
  };

  const newSchema = Object.assign(schema, {
    new: newFunc,
    jsonParse: parseJson,
    jsonParseSafe: parseJsonSafe,
    toPrompt: () => {
      // @ts-expect-error
      if (schema._def.typeName !== "ZodObject")
        throw new Error("Cannot use toPrompt on non-object schema");
      // @ts-expect-error
      return outputPrompt(schema);
    },
    stringify,
  });

  return newSchema;
};