Map over ZodObject and change keys to use a different schema.

[Kuirak]

Hello, I want to convert an existing schema (ZodObject) to use different schemas for the keys to realize using the same base schema for the config and the implementation (resolve some kind of references), here is an example:

// The origin schema
const statusSchema = z.object({
  largeTitle: z.boolean().default(false),
  status: z.string(),
  statusValue: z.number(),
})

// The target key type implementation
const sourceValueSchema = z.object({ sourceIndex: z.number(), key: z.string() });
const numberTypeSchema = z.object({ type: z.literal("number"), value: z.union([z.number(), sourceValueSchema]) });
const booleanTypeSchema = z.object({ type: z.literal("boolean"), value: z.union([z.boolean(), sourceValueSchema]) });
const stringTypeSchema = z.object({ type: z.literal("string"), value: z.union([z.string(), sourceValueSchema]) });

// The resulting Schema should look like this but I want to convert it programmatically:
const statusConfigSchema =  z.object({
  largeTitle: booleanTypeSchema,
  status: stringTypeSchema,
  statusValue: numberTypeSchema,
})

I tried to implement a convert function because I will have a lot of those and I want to avoid doing it manually, here is the function I came up with but I am losing the types:

const convertSchemaForConfig = <T extends ZodObject<K>, K extends ZodRawShape>(schema: T) => {
    type Keys = keyof typeof schema.shape;
    const shape: [Keys, ZodType][] = Object.entries(schema.shape);
    const newShape = Object.fromEntries(
      shape.map(([key, s]) => {
        if (s instanceof ZodString) {
          return [key, stringTypeSchema];
        } else if (s instanceof ZodNumber) {
          return [key, numberTypeSchema];
        } else if (s instanceof ZodBoolean) {
          return [key, booleanTypeSchema];
        }
        return [key, s];
      }),
    );
    return z.object(newShape);
};

Any Ideas how I could improve this?

[JacobWeisenburger]

Any Ideas how I could improve this?

A broad question like this is really hard to answer. If you are able to ask a more specific question or say what part you are wanting to improve, then it's much easier for someone else to answer.