jazz-tools 0.16.2 → 0.16.4

package/src/tools/implementation/zodSchema/schemaTypes/CoMapSchema.ts

@@ -5,24 +5,25 @@ import {
   DiscriminableCoValueSchemaDefinition,
   DiscriminableCoreCoValueSchema,
   Group,
-  PartialOnUndefined,
   RefsToResolve,
   RefsToResolveStrict,
   Resolved,
   Simplify,
   SubscribeListenerOptions,
+  coMapDefiner,
   coOptionalDefiner,
   hydrateCoreCoValueSchema,
   isAnyCoValueSchema,
 } from "../../../internal.js";
 import { AnonymousJazzAgent } from "../../anonymousJazzAgent.js";
 import { removeGetters } from "../../schemaUtils.js";
-import { CoFieldInit } from "../typeConverters/CoFieldInit.js";
+import { CoMapSchemaInit } from "../typeConverters/CoFieldInit.js";
 import { InstanceOrPrimitiveOfSchema } from "../typeConverters/InstanceOrPrimitiveOfSchema.js";
 import { InstanceOrPrimitiveOfSchemaCoValuesNullable } from "../typeConverters/InstanceOrPrimitiveOfSchemaCoValuesNullable.js";
 import { z } from "../zodReExport.js";
-import { AnyZodOrCoValueSchema } from "../zodSchema.js";
+import { AnyZodOrCoValueSchema, AnyZodSchema } from "../zodSchema.js";
 import { CoOptionalSchema } from "./CoOptionalSchema.js";
+import { CoreCoValueSchema } from "./CoValueSchema.js";
 
 export interface CoMapSchema<
   Shape extends z.core.$ZodLooseShape,
@@ -30,7 +31,7 @@ export interface CoMapSchema<
   Owner extends Account | Group = Account | Group,
 > extends CoreCoMapSchema<Shape, CatchAll> {
   create: (
-    init: Simplify<CoMapSchemaInit<Shape>>,
+    init: CoMapSchemaInit<Shape>,
     options?:
       | {
           owner: Owner;
@@ -134,6 +135,23 @@ export interface CoMapSchema<
   getCoValueClass: () => typeof CoMap;
 
   optional(): CoOptionalSchema<this>;
+
+  /**
+   * Creates a new CoMap schema by picking the specified keys from the original schema.
+   *
+   * @param keys - The keys to pick from the original schema.
+   * @returns A new CoMap schema with the picked keys.
+   */
+  pick<Keys extends keyof Shape>(
+    keys: { [key in Keys]: true },
+  ): CoMapSchema<Simplify<Pick<Shape, Keys>>, unknown, Owner>;
+
+  /**
+   * Creates a new CoMap schema by making all fields optional.
+   *
+   * @returns A new CoMap schema with all fields optional.
+   */
+  partial(): CoMapSchema<PartialShape<Shape>, CatchAll, Owner>;
 }
 
 export function createCoreCoMapSchema<
@@ -219,21 +237,44 @@ export function enrichCoMapSchema<
     getCoValueClass: () => {
       return coValueClass;
     },
-
     optional: () => {
       return coOptionalDefiner(coValueSchema);
     },
+    pick: <Keys extends keyof Shape>(keys: { [key in Keys]: true }) => {
+      const keysSet = new Set(Object.keys(keys));
+      const pickedShape: Record<string, AnyZodOrCoValueSchema> = {};
+
+      for (const [key, value] of Object.entries(coValueSchema.shape)) {
+        if (keysSet.has(key)) {
+          pickedShape[key] = value;
+        }
+      }
+
+      return coMapDefiner(pickedShape);
+    },
+    partial: () => {
+      const partialShape: Record<string, AnyZodOrCoValueSchema> = {};
+
+      for (const [key, value] of Object.entries(coValueSchema.shape)) {
+        if (isAnyCoValueSchema(value)) {
+          partialShape[key] = coOptionalDefiner(value);
+        } else {
+          partialShape[key] = z.optional(coValueSchema.shape[key]);
+        }
+      }
+
+      const partialCoMapSchema = coMapDefiner(partialShape);
+      if (coValueSchema.catchAll) {
+        return partialCoMapSchema.catchall(
+          coValueSchema.catchAll as unknown as AnyZodOrCoValueSchema,
+        );
+      }
+      return partialCoMapSchema;
+    },
   }) as unknown as CoMapSchema<Shape, CatchAll>;
   return coValueSchema;
 }
 
-// Due to a TS limitation with types that contain known properties and
-// an index signature, we cannot accept catchall properties on creation
-export type CoMapSchemaInit<Shape extends z.core.$ZodLooseShape> =
-  PartialOnUndefined<{
-    [key in keyof Shape]: CoFieldInit<Shape[key]>;
-  }>;
-
 export interface CoMapSchemaDefinition<
   Shape extends z.core.$ZodLooseShape = z.core.$ZodLooseShape,
   CatchAll extends AnyZodOrCoValueSchema | unknown = unknown,
@@ -270,3 +311,11 @@ export type CoMapInstanceCoValuesNullable<Shape extends z.core.$ZodLooseShape> =
       Shape[key]
     >;
   };
+
+export type PartialShape<Shape extends z.core.$ZodLooseShape> = Simplify<{
+  -readonly [key in keyof Shape]: Shape[key] extends AnyZodSchema
+    ? z.ZodOptional<Shape[key]>
+    : Shape[key] extends CoreCoValueSchema
+      ? CoOptionalSchema<Shape[key]>
+      : never;
+}>;

package/src/tools/implementation/zodSchema/typeConverters/CoFieldInit.ts

@@ -1,13 +1,78 @@
-import { NotNull } from "../../../internal.js";
+import {
+  CoDiscriminatedUnionSchema,
+  CoValueClass,
+  CoreCoFeedSchema,
+  CoreCoListSchema,
+  CoreCoMapSchema,
+  CoreCoRecordSchema,
+  CorePlainTextSchema,
+  PartialOnUndefined,
+  Simplify,
+} from "../../../internal.js";
+import { CoreCoOptionalSchema } from "../schemaTypes/CoOptionalSchema.js";
+import { CoreCoValueSchema } from "../schemaTypes/CoValueSchema.js";
+import { CoreRichTextSchema } from "../schemaTypes/RichTextSchema.js";
 import { z } from "../zodReExport.js";
-import { AnyZodOrCoValueSchema } from "../zodSchema.js";
-import { InstanceOrPrimitiveOfSchemaCoValuesNullable } from "./InstanceOrPrimitiveOfSchemaCoValuesNullable.js";
+import { AnyZodOrCoValueSchema, Loaded } from "../zodSchema.js";
+import { TypeOfZodSchema } from "./TypeOfZodSchema.js";
 
 /**
- * Returns the type of the value that should be used to initialize a coField
- * of the given schema.
+ * The type of value that can be used to initialize a CoField of the given schema.
+ *
+ * For CoValue fields, this can be either a shallowly-loaded CoValue instance
+ * or a JSON object that will be used to create the CoValue.
  */
-export type CoFieldInit<T extends AnyZodOrCoValueSchema> =
-  T extends z.core.$ZodNullable
-    ? InstanceOrPrimitiveOfSchemaCoValuesNullable<T>
-    : NotNull<InstanceOrPrimitiveOfSchemaCoValuesNullable<T>>;
+export type CoFieldInit<S extends CoValueClass | AnyZodOrCoValueSchema> =
+  S extends CoreCoValueSchema
+    ?
+        | Loaded<S>
+        | (S extends CoreCoRecordSchema<infer K, infer V>
+            ? CoMapSchemaInit<{ [key in z.output<K> & string]: V }>
+            : S extends CoreCoMapSchema<infer Shape>
+              ? CoMapSchemaInit<Shape>
+              : S extends CoreCoListSchema<infer T>
+                ? CoListInit<T>
+                : S extends CoreCoFeedSchema<infer T>
+                  ? CoFeedInit<T>
+                  : S extends CorePlainTextSchema | CoreRichTextSchema
+                    ? string
+                    : S extends CoreCoOptionalSchema<infer T>
+                      ? CoFieldInit<T> | undefined
+                      : S extends CoDiscriminatedUnionSchema<infer Members>
+                        ? CoFieldInit<Members[number]>
+                        : never)
+    : S extends z.core.$ZodType
+      ? TypeOfZodSchema<S>
+      : S extends CoValueClass
+        ? InstanceType<S>
+        : never;
+
+// Due to a TS limitation with types that contain known properties and
+// an index signature, we cannot accept catchall properties on creation
+export type CoMapSchemaInit<Shape extends z.core.$ZodLooseShape> = Simplify<
+  {
+    /**
+     * Cannot use {@link PartialOnUndefined} because evaluating CoFieldInit<Shape[Key]>
+     * to know if the value can be undefined does not work with recursive types.
+     */
+    [Key in keyof Shape as Shape[Key] extends
+      | CoreCoOptionalSchema
+      | z.core.$ZodOptional
+      ? never
+      : Key]: CoFieldInit<Shape[Key]>;
+  } & {
+    [Key in keyof Shape as Shape[Key] extends
+      | CoreCoOptionalSchema
+      | z.core.$ZodOptional
+      ? Key
+      : never]?: CoFieldInit<Shape[Key]>;
+  }
+>;
+
+export type CoListInit<T extends AnyZodOrCoValueSchema> = Simplify<
+  ReadonlyArray<CoFieldInit<T>>
+>;
+
+export type CoFeedInit<T extends AnyZodOrCoValueSchema> = Simplify<
+  ReadonlyArray<CoFieldInit<T>>
+>;

package/src/tools/implementation/zodSchema/typeConverters/InstanceOfSchema.ts

@@ -27,15 +27,20 @@ export type InstanceOfSchema<S extends CoValueClass | AnyZodOrCoValueSchema> =
   S extends CoreCoValueSchema
     ? S extends CoreAccountSchema<infer Shape>
       ? {
-          [key in keyof Shape]: InstanceOrPrimitiveOfSchema<Shape[key]>;
+          -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchema<
+            Shape[key]
+          >;
         } & Account
       : S extends CoreCoRecordSchema<infer K, infer V>
         ? {
-            [key in z.output<K> & string]: InstanceOrPrimitiveOfSchema<V>;
+            -readonly [key in z.output<K> &
+              string]: InstanceOrPrimitiveOfSchema<V>;
           } & CoMap
         : S extends CoreCoMapSchema<infer Shape, infer CatchAll>
           ? {
-              [key in keyof Shape]: InstanceOrPrimitiveOfSchema<Shape[key]>;
+              -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchema<
+                Shape[key]
+              >;
             } & (CatchAll extends AnyZodOrCoValueSchema
               ? {
                   [key: string]: InstanceOrPrimitiveOfSchema<CatchAll>;
@@ -53,7 +58,7 @@ export type InstanceOfSchema<S extends CoValueClass | AnyZodOrCoValueSchema> =
                   : S extends CoreFileStreamSchema
                     ? FileStream
                     : S extends CoreCoOptionalSchema<infer T>
-                      ? InstanceOrPrimitiveOfSchema<T>
+                      ? InstanceOrPrimitiveOfSchema<T> | undefined
                       : S extends CoDiscriminatedUnionSchema<infer Members>
                         ? InstanceOrPrimitiveOfSchema<Members[number]>
                         : never

package/src/tools/implementation/zodSchema/typeConverters/InstanceOfSchemaCoValuesNullable.ts

@@ -29,7 +29,7 @@ export type InstanceOfSchemaCoValuesNullable<
   ? S extends CoreAccountSchema<infer Shape>
     ?
         | ({
-            [key in keyof Shape]: InstanceOrPrimitiveOfSchemaCoValuesNullable<
+            -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchemaCoValuesNullable<
               Shape[key]
             >;
           } & Account)
@@ -37,14 +37,14 @@ export type InstanceOfSchemaCoValuesNullable<
     : S extends CoreCoRecordSchema<infer K, infer V>
       ?
           | ({
-              [key in z.output<K> &
+              -readonly [key in z.output<K> &
                 string]: InstanceOrPrimitiveOfSchemaCoValuesNullable<V>;
             } & CoMap)
           | null
       : S extends CoreCoMapSchema<infer Shape, infer CatchAll>
         ?
             | ({
-                [key in keyof Shape]: InstanceOrPrimitiveOfSchemaCoValuesNullable<
+                -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchemaCoValuesNullable<
                   Shape[key]
                 >;
               } & (CatchAll extends AnyZodOrCoValueSchema

package/src/tools/implementation/zodSchema/typeConverters/InstanceOrPrimitiveOfSchema.ts

@@ -1,122 +1,11 @@
-import { JsonValue } from "cojson";
 import {
-  Account,
   AnyZodOrCoValueSchema,
-  CoDiscriminatedUnionSchema,
-  CoFeed,
-  CoList,
-  CoMap,
-  CoPlainText,
-  CoRichText,
   CoValueClass,
-  CoreAccountSchema,
-  CoreCoRecordSchema,
-  FileStream,
-  Profile,
+  InstanceOfSchema,
 } from "../../../internal.js";
-import { CoreCoFeedSchema } from "../schemaTypes/CoFeedSchema.js";
-import { CoreCoListSchema } from "../schemaTypes/CoListSchema.js";
-import { CoreCoMapSchema } from "../schemaTypes/CoMapSchema.js";
-import { CoreCoOptionalSchema } from "../schemaTypes/CoOptionalSchema.js";
-import { CoreCoValueSchema } from "../schemaTypes/CoValueSchema.js";
-import { CoreFileStreamSchema } from "../schemaTypes/FileStreamSchema.js";
-import { CorePlainTextSchema } from "../schemaTypes/PlainTextSchema.js";
-import { CoreRichTextSchema } from "../schemaTypes/RichTextSchema.js";
 import { z } from "../zodReExport.js";
+import { TypeOfZodSchema } from "./TypeOfZodSchema.js";
 
 export type InstanceOrPrimitiveOfSchema<
   S extends CoValueClass | AnyZodOrCoValueSchema,
-> = S extends CoreCoValueSchema
-  ? S extends CoreAccountSchema<infer Shape>
-    ? {
-        -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchema<Shape[key]>;
-      } & { profile: Profile } & Account
-    : S extends CoreCoRecordSchema<infer K, infer V>
-      ? {
-          -readonly [key in z.output<K> &
-            string]: InstanceOrPrimitiveOfSchema<V>;
-        } & CoMap
-      : S extends CoreCoMapSchema<infer Shape, infer CatchAll>
-        ? {
-            -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchema<
-              Shape[key]
-            >;
-          } & (CatchAll extends AnyZodOrCoValueSchema
-            ? {
-                [key: string]: InstanceOrPrimitiveOfSchema<CatchAll>;
-              }
-            : {}) &
-            CoMap
-        : S extends CoreCoListSchema<infer T>
-          ? CoList<InstanceOrPrimitiveOfSchema<T>>
-          : S extends CoreCoFeedSchema<infer T>
-            ? CoFeed<InstanceOrPrimitiveOfSchema<T>>
-            : S extends CorePlainTextSchema
-              ? CoPlainText
-              : S extends CoreRichTextSchema
-                ? CoRichText
-                : S extends CoreFileStreamSchema
-                  ? FileStream
-                  : S extends CoreCoOptionalSchema<infer T>
-                    ? InstanceOrPrimitiveOfSchema<T> | undefined
-                    : S extends CoDiscriminatedUnionSchema<infer Members>
-                      ? InstanceOrPrimitiveOfSchema<Members[number]>
-                      : never
-  : S extends z.core.$ZodType
-    ? S extends z.core.$ZodOptional<infer Inner extends z.core.$ZodType>
-      ? InstanceOrPrimitiveOfSchema<Inner> | undefined
-      : S extends z.core.$ZodNullable<infer Inner extends z.core.$ZodType>
-        ? InstanceOrPrimitiveOfSchema<Inner> | null
-        : S extends z.ZodJSONSchema
-          ? JsonValue
-          : S extends z.core.$ZodUnion<infer Members extends z.core.$ZodType[]>
-            ? InstanceOrPrimitiveOfSchema<Members[number]>
-            : // primitives below here - we manually traverse to ensure we only allow what we can handle
-              S extends z.core.$ZodObject<infer Shape>
-              ? {
-                  -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchema<
-                    Shape[key]
-                  >;
-                }
-              : S extends z.core.$ZodArray<infer Item extends z.core.$ZodType>
-                ? InstanceOrPrimitiveOfSchema<Item>[]
-                : S extends z.core.$ZodTuple<
-                      infer Items extends readonly z.core.$ZodType[]
-                    >
-                  ? {
-                      [key in keyof Items]: InstanceOrPrimitiveOfSchema<
-                        Items[key]
-                      >;
-                    }
-                  : S extends z.core.$ZodString
-                    ? string
-                    : S extends z.core.$ZodNumber
-                      ? number
-                      : S extends z.core.$ZodBoolean
-                        ? boolean
-                        : S extends z.core.$ZodLiteral<infer Literal>
-                          ? Literal
-                          : S extends z.core.$ZodDate
-                            ? Date
-                            : S extends z.core.$ZodEnum<infer Enum>
-                              ? Enum[keyof Enum]
-                              : S extends z.core.$ZodTemplateLiteral<
-                                    infer pattern
-                                  >
-                                ? pattern
-                                : S extends z.core.$ZodReadonly<
-                                      infer Inner extends z.core.$ZodType
-                                    >
-                                  ? InstanceOrPrimitiveOfSchema<Inner>
-                                  : S extends z.core.$ZodDefault<
-                                        infer Default extends z.core.$ZodType
-                                      >
-                                    ? InstanceOrPrimitiveOfSchema<Default>
-                                    : S extends z.core.$ZodCatch<
-                                          infer Catch extends z.core.$ZodType
-                                        >
-                                      ? InstanceOrPrimitiveOfSchema<Catch>
-                                      : never
-    : S extends CoValueClass
-      ? InstanceType<S>
-      : never;
+> = S extends z.core.$ZodType ? TypeOfZodSchema<S> : InstanceOfSchema<S>;

package/src/tools/implementation/zodSchema/typeConverters/InstanceOrPrimitiveOfSchemaCoValuesNullable.ts

@@ -1,137 +1,13 @@
-import { JsonValue } from "cojson";
 import {
-  Account,
   AnyZodOrCoValueSchema,
-  CoDiscriminatedUnionSchema,
-  CoFeed,
-  CoList,
-  CoMap,
-  CoPlainText,
-  CoRichText,
   CoValueClass,
-  CoreAccountSchema,
-  CoreCoRecordSchema,
-  FileStream,
-  InstanceOrPrimitiveOfSchema,
-  Profile,
+  InstanceOfSchemaCoValuesNullable,
 } from "../../../internal.js";
-import { CoreCoFeedSchema } from "../schemaTypes/CoFeedSchema.js";
-import { CoreCoListSchema } from "../schemaTypes/CoListSchema.js";
-import { CoreCoMapSchema } from "../schemaTypes/CoMapSchema.js";
-import { CoreCoOptionalSchema } from "../schemaTypes/CoOptionalSchema.js";
-import { CoreCoValueSchema } from "../schemaTypes/CoValueSchema.js";
-import { CoreFileStreamSchema } from "../schemaTypes/FileStreamSchema.js";
-import { CorePlainTextSchema } from "../schemaTypes/PlainTextSchema.js";
-import { CoreRichTextSchema } from "../schemaTypes/RichTextSchema.js";
 import { z } from "../zodReExport.js";
+import { TypeOfZodSchema } from "./TypeOfZodSchema.js";
 
 export type InstanceOrPrimitiveOfSchemaCoValuesNullable<
   S extends CoValueClass | AnyZodOrCoValueSchema,
-> = S extends CoreCoValueSchema
-  ? S extends CoreAccountSchema<infer Shape>
-    ?
-        | ({
-            -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchemaCoValuesNullable<
-              Shape[key]
-            >;
-          } & { profile: Profile | null } & Account)
-        | null
-    : S extends CoreCoRecordSchema<infer K, infer V>
-      ?
-          | ({
-              -readonly [key in z.output<K> &
-                string]: InstanceOrPrimitiveOfSchemaCoValuesNullable<V>;
-            } & CoMap)
-          | null
-      : S extends CoreCoMapSchema<infer Shape, infer CatchAll>
-        ?
-            | ({
-                -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchemaCoValuesNullable<
-                  Shape[key]
-                >;
-              } & (CatchAll extends AnyZodOrCoValueSchema
-                ? {
-                    [
-                      key: string
-                    ]: InstanceOrPrimitiveOfSchemaCoValuesNullable<CatchAll>;
-                  }
-                : {}) &
-                CoMap)
-            | null
-        : S extends CoreCoListSchema<infer T>
-          ? CoList<InstanceOrPrimitiveOfSchemaCoValuesNullable<T>> | null
-          : S extends CoreCoFeedSchema<infer T>
-            ? CoFeed<InstanceOrPrimitiveOfSchemaCoValuesNullable<T>> | null
-            : S extends CorePlainTextSchema
-              ? CoPlainText | null
-              : S extends CoreRichTextSchema
-                ? CoRichText | null
-                : S extends CoreFileStreamSchema
-                  ? FileStream | null
-                  : S extends CoreCoOptionalSchema<infer T>
-                    ? InstanceOrPrimitiveOfSchemaCoValuesNullable<T> | undefined
-                    : S extends CoDiscriminatedUnionSchema<infer Members>
-                      ? InstanceOrPrimitiveOfSchemaCoValuesNullable<
-                          Members[number]
-                        >
-                      : never
-  : S extends z.core.$ZodType
-    ? S extends z.core.$ZodOptional<infer Inner extends z.core.$ZodType>
-      ? InstanceOrPrimitiveOfSchemaCoValuesNullable<Inner> | undefined
-      : S extends z.core.$ZodNullable<infer Inner extends z.core.$ZodType>
-        ? InstanceOrPrimitiveOfSchemaCoValuesNullable<Inner> | null
-        : S extends z.ZodJSONSchema
-          ? JsonValue
-          : S extends z.core.$ZodUnion<
-                infer Members extends readonly z.core.$ZodType[]
-              >
-            ? InstanceOrPrimitiveOfSchemaCoValuesNullable<Members[number]>
-            : // primitives below here - we manually traverse to ensure we only allow what we can handle
-              S extends z.core.$ZodObject<infer Shape>
-              ? {
-                  -readonly [key in keyof Shape]: InstanceOrPrimitiveOfSchema<
-                    Shape[key]
-                  >;
-                }
-              : S extends z.core.$ZodArray<infer Item extends z.core.$ZodType>
-                ? InstanceOrPrimitiveOfSchema<Item>[]
-                : S extends z.core.$ZodTuple<
-                      infer Items extends z.core.$ZodType[]
-                    >
-                  ? {
-                      [key in keyof Items]: InstanceOrPrimitiveOfSchema<
-                        Items[key]
-                      >;
-                    }
-                  : S extends z.core.$ZodString
-                    ? string
-                    : S extends z.core.$ZodNumber
-                      ? number
-                      : S extends z.core.$ZodBoolean
-                        ? boolean
-                        : S extends z.core.$ZodLiteral<infer Literal>
-                          ? Literal
-                          : S extends z.core.$ZodDate
-                            ? Date
-                            : S extends z.core.$ZodEnum<infer Enum>
-                              ? Enum[keyof Enum]
-                              : S extends z.core.$ZodTemplateLiteral<
-                                    infer pattern
-                                  >
-                                ? pattern
-                                : S extends z.core.$ZodReadonly<
-                                      infer Inner extends z.core.$ZodType
-                                    >
-                                  ? InstanceOrPrimitiveOfSchema<Inner>
-                                  : S extends z.core.$ZodDefault<
-                                        infer Default extends z.core.$ZodType
-                                      >
-                                    ? InstanceOrPrimitiveOfSchema<Default>
-                                    : S extends z.core.$ZodCatch<
-                                          infer Catch extends z.core.$ZodType
-                                        >
-                                      ? InstanceOrPrimitiveOfSchema<Catch>
-                                      : never
-    : S extends CoValueClass
-      ? InstanceType<S> | null
-      : never;
+> = S extends z.core.$ZodType
+  ? TypeOfZodSchema<S>
+  : InstanceOfSchemaCoValuesNullable<S>;

package/src/tools/implementation/zodSchema/typeConverters/TypeOfZodSchema.ts

@@ -0,0 +1,79 @@
+import { JsonValue } from "cojson";
+import { PartialOnUndefined } from "../../../internal.js";
+import { z } from "../zodReExport.js";
+
+// Copied from https://github.com/colinhacks/zod/blob/7e7e3461aceecf3633e158df50d6bc852e7cdf45/packages/zod/src/v4/core/schemas.ts#L1591,
+// since this type is not exported by Zod
+type OptionalInSchema = {
+  _zod: {
+    optin: "optional";
+  };
+};
+
+/**
+ * Get type from Zod schema definition.
+ *
+ * Similar to `z.infer`, but we manually traverse Zod types to ensure we only allow what we can handle
+ */
+export type TypeOfZodSchema<S extends z.core.$ZodType> =
+  S extends z.core.$ZodOptional<infer Inner extends z.core.$ZodType>
+    ? TypeOfZodSchema<Inner> | undefined
+    : S extends z.core.$ZodNullable<infer Inner extends z.core.$ZodType>
+      ? TypeOfZodSchema<Inner> | null
+      : S extends z.ZodJSONSchema
+        ? JsonValue
+        : S extends z.core.$ZodUnion<
+              infer Members extends readonly z.core.$ZodType[]
+            >
+          ? TypeOfZodSchema<Members[number]>
+          : S extends z.core.$ZodObject<infer Shape>
+            ? /**
+               * Cannot use {@link PartialOnUndefined} because evaluating TypeOfZodSchema<Shape[key]>
+               * to know if the value can be undefined does not work with recursive types.
+               */
+              {
+                -readonly [key in keyof Shape as Shape[key] extends OptionalInSchema
+                  ? never
+                  : key]: TypeOfZodSchema<Shape[key]>;
+              } & {
+                -readonly [key in keyof Shape as Shape[key] extends OptionalInSchema
+                  ? key
+                  : never]?: TypeOfZodSchema<Shape[key]>;
+              }
+            : S extends z.core.$ZodArray<infer Item extends z.core.$ZodType>
+              ? TypeOfZodSchema<Item>[]
+              : S extends z.core.$ZodTuple<
+                    infer Items extends readonly z.core.$ZodType[]
+                  >
+                ? {
+                    [key in keyof Items]: TypeOfZodSchema<Items[key]>;
+                  }
+                : S extends z.core.$ZodString
+                  ? string
+                  : S extends z.core.$ZodNumber
+                    ? number
+                    : S extends z.core.$ZodBoolean
+                      ? boolean
+                      : S extends z.core.$ZodLiteral<infer Literal>
+                        ? Literal
+                        : S extends z.core.$ZodDate
+                          ? Date
+                          : S extends z.core.$ZodEnum<infer Enum>
+                            ? Enum[keyof Enum]
+                            : S extends z.core.$ZodTemplateLiteral<
+                                  infer pattern
+                                >
+                              ? pattern
+                              : S extends z.core.$ZodReadonly<
+                                    infer Inner extends z.core.$ZodType
+                                  >
+                                ? TypeOfZodSchema<Inner>
+                                : S extends z.core.$ZodDefault<
+                                      infer Default extends z.core.$ZodType
+                                    >
+                                  ? TypeOfZodSchema<Default>
+                                  : S extends z.core.$ZodCatch<
+                                        infer Catch extends z.core.$ZodType
+                                      >
+                                    ? TypeOfZodSchema<Catch>
+                                    : never;

package/src/tools/implementation/zodSchema/unionUtils.ts

@@ -7,6 +7,7 @@ import {
   CoreCoMapSchema,
   DiscriminableCoValueSchemas,
   DiscriminableCoreCoValueSchema,
+  SchemaUnionDiscriminator,
 } from "../../internal.js";
 import {
   hydrateCoreCoValueSchema,
@@ -56,21 +57,17 @@ export function schemaUnionDiscriminatorFor(
       }
     }
 
-    const determineSchema = (_raw: RawCoMap | RawAccount | RawCoList) => {
-      if (_raw instanceof RawCoList) {
-        throw new Error(
-          "co.discriminatedUnion() of collaborative types is not supported for CoLists",
-        );
-      }
-
+    const determineSchema: SchemaUnionDiscriminator<CoMap> = (
+      discriminable,
+    ) => {
       for (const option of availableOptions) {
         let match = true;
 
         for (const key of Object.keys(discriminatorMap)) {
           const discriminatorDef = (option as CoreCoMapSchema).getDefinition()
-            .shape[key as string];
+            .shape[key];
 
-          const discriminatorValue = (_raw as RawCoMap).get(key as string);
+          const discriminatorValue = discriminable.get(key);
 
           if (discriminatorValue && typeof discriminatorValue === "object") {
             throw new Error("Discriminator must be a primitive value");