ZodOptional always return unknown, if it's not nested

[bikoevD]

I'm trying to add a handler for the zod schema. If the scheme contains an optional value, then it infer the unknown in the types, but if the optional value is nested, for example, inside another object, then it works correctly.

import { ObjectId } from "mongodb";
import { z } from "zod";

type ZodType = z.ZodType<any, any, any>;

/**
 * Transforms a Zod schema to handle MongoDB ObjectId fields
 * @param schema The Zod schema to transform
 * @returns A new schema with transformed ID fields
 */
function createMongoSchema<T extends z.ZodObject<any, any, any>>(
  schema: T
): z.ZodObject<
  {
    [k in keyof z.infer<T> as k extends "id" ? "_id" : k]: k extends "id"
      ? z.ZodType<ObjectId>
      : k extends `${string}Id` | `${string}Ids`
        ? z.infer<T>[k] extends any[]
          ? z.ZodArray<z.ZodType<ObjectId>>
          : z.ZodType<ObjectId>
        : T["shape"][k] extends z.ZodOptional<infer U>
          ? z.ZodOptional<U extends z.ZodObject<any, any, any> ? ReturnType<typeof createMongoSchema<U>> : U>
          : T["shape"][k];
  },
  "strip"
> {
  if (!(schema instanceof z.ZodType)) {
    return schema;
  }

  if (schema instanceof z.ZodObject) {
    const shape = schema._def.shape();
    const newShape: Record<string, ZodType> = {};
    for (const [key, value] of Object.entries(shape)) {
      if (isZodId(value as ZodType)) {
        // Handle 'id' field specially
        if (key === "id") {
          newShape._id = z.instanceof(ObjectId);
          continue;
        }
        // Handle other ID fields
        newShape[key] = z.instanceof(ObjectId);
        continue;
      }

      // Recursively transform nested objects and arrays
      if (value instanceof z.ZodObject) {
        newShape[key] = createMongoSchema(value);
      } else if (value instanceof z.ZodArray) {
        const elementSchema = value.element;
        if (elementSchema instanceof z.ZodObject) {
          newShape[key] = z.array(createMongoSchema(elementSchema));
        } else if (isZodId(elementSchema)) {
          newShape[key] = z.array(z.instanceof(ObjectId));
        } else {
          newShape[key] = value;
        }
      } else if (value instanceof z.ZodOptional) {
        const innerSchema = value.unwrap();
        if (innerSchema instanceof z.ZodObject) {
          newShape[key] = z.optional(createMongoSchema(innerSchema));
        } else if (isZodId(innerSchema)) {
          newShape[key] = z.optional(z.instanceof(ObjectId));
        } else {
          newShape[key] = z.optional(innerSchema);
        }
      } else {
        newShape[key] = value as ZodType;
      }
    }

    return z.object(newShape) as any;
  }
  return schema;
}

/**
 * Checks if a Zod type is created using zodId()
 */
function isZodId(schema: ZodType): boolean {
  return (
    schema instanceof z.ZodString &&
    schema._def.checks.some((check) => check.kind === "regex" && check.regex.source === "^[a-f\\d]{24}$")
  );
}

export { createMongoSchema };

const collectionSchema = z.object({
  id: zodId(),
  commonProps: z.object({
    arg1: z.boolean().optional(),
    arg2: z.string().optional(),
    arg3: z.number()
  }),
  commonPropsOptional: z
    .object({
      arg1: z.boolean().optional(),
      arg2: z.string().optional(),
      arg3: z.number()
    })
    .optional(),
  stringValueIds: z.array(zodId()),
  stringValueOptional: z.string().optional()
});

const mongoSchema = createMongoSchema(collectionSchema);
type SchemaType = z.infer<typeof mongoSchema>;

Any ideas how to fix this?

[bikoevD]

I fixed the issues with unknown, now it's working correctly. But the types turned out to be quite massive, maybe there is a way to reduce them?

type MongoSchema<
  T extends z.ZodRawShape,
  UnknownKeys extends z.UnknownKeysParam,
  Catchall extends z.ZodTypeAny
> = z.ZodObject<
  {
    [K in keyof T as K extends "id" ? "_id" : K]: K extends "id"
      ? z.ZodType<ObjectId>
      : K extends `${string}Id` | `${string}Ids`
        ? T[K] extends z.ZodOptional<z.ZodArray<any>>
          ? z.ZodOptional<z.ZodArray<z.ZodType<ObjectId>>>
          : T[K] extends z.ZodOptional<any>
            ? z.ZodOptional<z.ZodType<ObjectId>>
            : T[K] extends z.ZodArray<any>
              ? z.ZodArray<z.ZodType<ObjectId>>
              : z.ZodType<ObjectId>
        : T[K] extends z.ZodOptional<infer U>
          ? z.ZodOptional<TransformType<U>>
          : TransformType<T[K]>;
  },
  UnknownKeys,
  Catchall
>;

// Helper type to recursively transform nested objects and arrays
type TransformType<T> =
  T extends z.ZodArray<infer W>
    ? W extends z.ZodType<any>
      ? z.ZodArray<TransformType<W>>
      : T
    : T extends z.ZodObject<infer U extends z.ZodRawShape, any, any>
      ? MongoSchema<U, z.UnknownKeysParam, z.ZodTypeAny>
      : T;

/**
 * Transforms a Zod schema to handle MongoDB ObjectId fields
 * @param schema The Zod schema to transform
 * @returns A new schema with transformed ID fields
 */
function createMongoSchema<
  T extends z.ZodRawShape,
  UnknownKeys extends z.UnknownKeysParam,
  Catchall extends z.ZodTypeAny
>(schema: z.ZodObject<T, UnknownKeys, Catchall>): MongoSchema<T, UnknownKeys, Catchall> {
...