jazz-tools 0.15.10 → 0.15.12

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

@@ -6,9 +6,12 @@ import {
 } from "../../../internal.js";
 import { z } from "../zodReExport.js";
 
-export type FileStreamSchema = z.core.$ZodCustom<FileStream, unknown> & {
+export type AnyFileStreamSchema = z.core.$ZodCustom<FileStream, unknown> & {
   collaborative: true;
   builtin: "FileStream";
+};
+
+export type FileStreamSchema = AnyFileStreamSchema & {
   create(options?: { owner?: Account | Group } | Account | Group): FileStream;
   createFromBlob(
     blob: Blob | File,
@@ -40,5 +43,32 @@ export type FileStreamSchema = z.core.$ZodCustom<FileStream, unknown> & {
     id: string,
     listener: (value: FileStream, unsubscribe: () => void) => void,
   ): () => void;
-  getCoSchema: () => typeof FileStream;
+  getCoValueClass: () => typeof FileStream;
 };
+
+export function enrichFileStreamSchema(
+  schema: AnyFileStreamSchema,
+  coValueClass: typeof FileStream,
+): FileStreamSchema {
+  return Object.assign(schema, {
+    create: (...args: any[]) => {
+      return coValueClass.create(...args);
+    },
+    createFromBlob: (...args: [any, ...any[]]) => {
+      return coValueClass.createFromBlob(...args);
+    },
+    load: (...args: [any, ...any[]]) => {
+      return coValueClass.load(...args);
+    },
+    loadAsBlob: (...args: [any, ...any[]]) => {
+      return coValueClass.loadAsBlob(...args);
+    },
+    subscribe: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as FileStreamSchema;
+}

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

@@ -3,9 +3,12 @@ import { Account, CoPlainText, Group } from "../../../internal.js";
 import { AnonymousJazzAgent } from "../../anonymousJazzAgent.js";
 import { z } from "../zodReExport.js";
 
-export type PlainTextSchema = z.core.$ZodCustom<CoPlainText, unknown> & {
+export type AnyPlainTextSchema = z.core.$ZodCustom<CoPlainText, unknown> & {
   collaborative: true;
   builtin: "CoPlainText";
+};
+
+export type PlainTextSchema = AnyPlainTextSchema & {
   create(
     text: string,
     options?: { owner: Account | Group } | Account | Group,
@@ -24,5 +27,30 @@ export type PlainTextSchema = z.core.$ZodCustom<CoPlainText, unknown> & {
     listener: (value: CoPlainText, unsubscribe: () => void) => void,
   ): () => void;
   fromRaw(raw: RawCoPlainText): CoPlainText;
-  getCoSchema: () => typeof CoPlainText;
+  getCoValueClass: () => typeof CoPlainText;
 };
+
+export function enrichPlainTextSchema(
+  schema: AnyPlainTextSchema,
+  coValueClass: typeof CoPlainText,
+): PlainTextSchema {
+  return Object.assign(schema, {
+    create: (...args: [any, ...any[]]) => {
+      return coValueClass.create(...args);
+    },
+    load: (...args: [any, ...any[]]) => {
+      return coValueClass.load(...args);
+    },
+    subscribe: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    fromRaw: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.fromRaw(...args);
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as PlainTextSchema;
+}

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

@@ -2,9 +2,12 @@ import { Account, CoRichText, Group } from "../../../internal.js";
 import { AnonymousJazzAgent } from "../../anonymousJazzAgent.js";
 import { z } from "../zodReExport.js";
 
-export type RichTextSchema = z.core.$ZodCustom<CoRichText, unknown> & {
+export type AnyRichTextSchema = z.core.$ZodCustom<CoRichText, unknown> & {
   collaborative: true;
   builtin: "CoRichText";
+};
+
+export type RichTextSchema = AnyRichTextSchema & {
   create(
     text: string,
     options?: { owner: Account | Group } | Account | Group,
@@ -22,5 +25,26 @@ export type RichTextSchema = z.core.$ZodCustom<CoRichText, unknown> & {
     id: string,
     listener: (value: CoRichText, unsubscribe: () => void) => void,
   ): () => void;
-  getCoSchema: () => typeof CoRichText;
+  getCoValueClass: () => typeof CoRichText;
 };
+
+export function enrichRichTextSchema(
+  schema: AnyRichTextSchema,
+  coValueClass: typeof CoRichText,
+): RichTextSchema {
+  return Object.assign(schema, {
+    create: (...args: [any, ...any[]]) => {
+      return coValueClass.create(...args);
+    },
+    load: (...args: [any, ...any[]]) => {
+      return coValueClass.load(...args);
+    },
+    subscribe: (...args: [any, ...any[]]) => {
+      // @ts-expect-error
+      return coValueClass.subscribe(...args);
+    },
+    getCoValueClass: () => {
+      return coValueClass;
+    },
+  }) as unknown as RichTextSchema;
+}

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

@@ -1,5 +1,7 @@
 import {
   Account,
+  AnyAccountSchema,
+  AnyCoRecordSchema,
   CoFeed,
   CoList,
   CoMap,
@@ -11,24 +13,19 @@ import {
 import { AnyCoFeedSchema } from "../schemaTypes/CoFeedSchema.js";
 import { AnyCoListSchema } from "../schemaTypes/CoListSchema.js";
 import { AnyCoMapSchema } from "../schemaTypes/CoMapSchema.js";
-import { FileStreamSchema } from "../schemaTypes/FileStreamSchema.js";
-import { PlainTextSchema } from "../schemaTypes/PlainTextSchema.js";
-import { RichTextSchema } from "../schemaTypes/RichTextSchema.js";
+import { AnyFileStreamSchema } from "../schemaTypes/FileStreamSchema.js";
+import { AnyPlainTextSchema } from "../schemaTypes/PlainTextSchema.js";
+import { AnyRichTextSchema } from "../schemaTypes/RichTextSchema.js";
 import { z } from "../zodReExport.js";
 import { InstanceOrPrimitiveOfSchema } from "./InstanceOrPrimitiveOfSchema.js";
 
 export type InstanceOfSchema<S extends CoValueClass | z.core.$ZodType> =
   S extends z.core.$ZodType
-    ? S extends z.core.$ZodObject<infer Shape> & {
-        collaborative: true;
-        builtin: "Account";
-      }
+    ? S extends AnyAccountSchema<infer Shape>
       ? {
           [key in keyof Shape]: InstanceOrPrimitiveOfSchema<Shape[key]>;
         } & Account
-      : S extends z.core.$ZodRecord<infer K, infer V> & {
-            collaborative: true;
-          }
+      : S extends AnyCoRecordSchema<infer K, infer V>
         ? {
             [key in z.output<K> & string]: InstanceOrPrimitiveOfSchema<V>;
           } & CoMap
@@ -45,11 +42,11 @@ export type InstanceOfSchema<S extends CoValueClass | z.core.$ZodType> =
             ? CoList<InstanceOrPrimitiveOfSchema<T>>
             : S extends AnyCoFeedSchema<infer T>
               ? CoFeed<InstanceOrPrimitiveOfSchema<T>>
-              : S extends PlainTextSchema
+              : S extends AnyPlainTextSchema
                 ? CoPlainText
-                : S extends RichTextSchema
+                : S extends AnyRichTextSchema
                   ? CoRichText
-                  : S extends FileStreamSchema
+                  : S extends AnyFileStreamSchema
                     ? FileStream
                     : S extends z.core.$ZodOptional<infer Inner>
                       ? InstanceOrPrimitiveOfSchema<Inner>

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

@@ -1,5 +1,12 @@
 import { RawAccount, RawCoList, RawCoMap } from "cojson";
-import { zodSchemaToCoSchema } from "./runtimeConverters/zodSchemaToCoSchema.js";
+import { CoValueClass, CoValueFromRaw } from "../../internal.js";
+import {
+  isAnyCoValueSchema,
+  zodSchemaToCoSchema,
+} from "./runtimeConverters/zodSchemaToCoSchema.js";
+import { AccountSchema } from "./schemaTypes/AccountSchema.js";
+import { CoListSchema } from "./schemaTypes/CoListSchema.js";
+import { CoMapSchema } from "./schemaTypes/CoMapSchema.js";
 import { z } from "./zodReExport.js";
 
 export function schemaUnionDiscriminatorFor(
@@ -8,7 +15,7 @@ export function schemaUnionDiscriminatorFor(
   if (isUnionOfCoMapsDeeply(schema)) {
     if (!schema._zod.disc || schema._zod.disc.size == 0) {
       throw new Error(
-        "z.union() of collaborative types is not supported, use z.discriminatedUnion() instead",
+        "z.union() of collaborative types is not supported, use co.discriminatedUnion() instead",
       );
     }
 
@@ -17,14 +24,14 @@ export function schemaUnionDiscriminatorFor(
 
     if (!field) {
       throw new Error(
-        "z.discriminatedUnion() of collaborative types with non-existent discriminator key is not supported",
+        "co.discriminatedUnion() of collaborative types with non-existent discriminator key is not supported",
       );
     }
 
     for (const value of field.values) {
       if (typeof value !== "string" && typeof value !== "number") {
         throw new Error(
-          "z.discriminatedUnion() of collaborative types with non-string or non-number discriminator value is not supported",
+          "co.discriminatedUnion() of collaborative types with non-string or non-number discriminator value is not supported",
         );
       }
     }
@@ -42,7 +49,7 @@ export function schemaUnionDiscriminatorFor(
         }
       } else {
         throw new Error(
-          "Unsupported zod type in z.discriminatedUnion() of collaborative types",
+          "Unsupported zod type in co.discriminatedUnion() of collaborative types",
         );
       }
     }
@@ -50,7 +57,7 @@ export function schemaUnionDiscriminatorFor(
     const determineSchema = (_raw: RawCoMap | RawAccount | RawCoList) => {
       if (_raw instanceof RawCoList) {
         throw new Error(
-          "z.discriminatedUnion() of collaborative types is not supported for CoLists",
+          "co.discriminatedUnion() of collaborative types is not supported for CoLists",
         );
       }
 
@@ -90,12 +97,17 @@ export function schemaUnionDiscriminatorFor(
         }
 
         if (match) {
-          return zodSchemaToCoSchema(option);
+          const coValueSchema = zodSchemaToCoSchema(option) as
+            | CoMapSchema<any>
+            | AccountSchema
+            | CoListSchema<any>;
+          return coValueSchema.getCoValueClass() as CoValueClass<any> &
+            CoValueFromRaw<any>;
         }
       }
 
       throw new Error(
-        "z.discriminatedUnion() of collaborative types with no matching discriminator value found",
+        "co.discriminatedUnion() of collaborative types with no matching discriminator value found",
       );
     };
 
@@ -120,11 +132,7 @@ export function isUnionOfCoMapsDeeply(
 function isCoMapOrUnionOfCoMapsDeeply(
   schema: z.core.$ZodType,
 ): schema is z.core.$ZodDiscriminatedUnion {
-  if (
-    schema instanceof z.core.$ZodObject &&
-    "collaborative" in schema &&
-    schema.collaborative
-  ) {
+  if (schema instanceof z.core.$ZodObject && isAnyCoValueSchema(schema)) {
     return true;
   } else if (schema instanceof z.core.$ZodUnion) {
     return schema._zod.def.options.every(isCoMapOrUnionOfCoMapsDeeply);
@@ -137,6 +145,6 @@ export function isUnionOfPrimitivesDeeply(schema: z.core.$ZodType) {
   if (schema instanceof z.core.$ZodUnion) {
     return schema._zod.def.options.every(isUnionOfPrimitivesDeeply);
   } else {
-    return !("collaborative" in schema);
+    return !isAnyCoValueSchema(schema);
   }
 }

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

@@ -1,8 +1,11 @@
 import {
-  type Account,
-  type AccountCreationProps,
   type AccountSchema,
-  type AnyCoMapSchema,
+  AnyAccountSchema,
+  AnyCoFeedSchema,
+  AnyCoListSchema,
+  AnyCoSchema,
+  AnyFileStreamSchema,
+  AnyPlainTextSchema,
   BaseAccountShape,
   CoFeed,
   type CoFeedSchema,
@@ -20,121 +23,31 @@ import {
   type Simplify,
   zodSchemaToCoSchema,
 } from "../../internal.js";
-import { RichTextSchema } from "./schemaTypes/RichTextSchema.js";
+import {
+  AnyDiscriminableCoSchema,
+  CoDiscriminatedUnionSchema,
+} from "./schemaTypes/CoDiscriminatedUnionSchema.js";
+import {
+  CoOptionalSchema,
+  createCoOptionalSchema,
+} from "./schemaTypes/CoOptionalSchema.js";
+import {
+  AnyRichTextSchema,
+  RichTextSchema,
+} from "./schemaTypes/RichTextSchema.js";
 import { z } from "./zodReExport.js";
 
-function enrichCoMapSchema<Shape extends z.core.$ZodLooseShape>(
-  schema: z.ZodObject<
-    { -readonly [P in keyof Shape]: Shape[P] },
-    z.core.$strip
-  >,
-) {
-  const baseCatchall = schema.catchall;
-
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    findUnique: (...args: any[]) => {
-      return coSchema.findUnique(...args);
-    },
-    upsertUnique: (...args: any[]) => {
-      return coSchema.upsertUnique(...args);
-    },
-    loadUnique: (...args: any[]) => {
-      return coSchema.loadUnique(...args);
-    },
-    catchall: (index: z.core.$ZodType) => {
-      return enrichCoMapSchema(baseCatchall(index));
-    },
-    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
-      return Object.assign(schema, helpers(schema));
-    },
-    withMigration: (migration: (value: any) => undefined) => {
-      coSchema.prototype.migrate = migration;
-
-      return enrichedSchema;
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as CoMapSchema<Shape>;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
 export const coMapDefiner = <Shape extends z.core.$ZodLooseShape>(
   shape: Shape,
 ): CoMapSchema<Shape> => {
   const objectSchema = z.object(shape).meta({
     collaborative: true,
   });
-
-  return enrichCoMapSchema(objectSchema);
-};
-
-function enrichAccountSchema<Shape extends BaseAccountShape>(
-  schema: z.ZodObject<Shape, z.core.$strip>,
-) {
-  const enrichedSchema = Object.assign(schema, {
+  const enrichedSchema = Object.assign(objectSchema, {
     collaborative: true,
-    builtin: "Account",
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    createAs: (...args: any[]) => {
-      return coSchema.createAs(...args);
-    },
-    getMe: (...args: any[]) => {
-      return coSchema.getMe(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
-      return Object.assign(schema, helpers(schema));
-    },
-    fromRaw: (...args: any[]) => {
-      return coSchema.fromRaw(...args);
-    },
-    withMigration: (
-      migration: (
-        value: any,
-        creationProps?: AccountCreationProps,
-      ) => void | Promise<void>,
-    ) => {
-      (coSchema.prototype as Account).migrate = async function (
-        this,
-        creationProps,
-      ) {
-        await migration(this, creationProps);
-      };
-
-      return enrichedSchema;
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as AccountSchema<Shape>;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
+  });
+  return zodSchemaToCoSchema(enrichedSchema);
+};
 
 /**
  * Defines a collaborative account schema for Jazz applications.
@@ -183,11 +96,14 @@ export const coAccountDefiner = <Shape extends BaseAccountShape>(
     root: coMapDefiner({}),
   } as unknown as Shape,
 ): AccountSchema<Shape> => {
-  const objectSchema = z.object(shape).meta({
+  const schema = z.object(shape).meta({
     collaborative: true,
   });
-
-  return enrichAccountSchema(objectSchema) as unknown as AccountSchema<Shape>;
+  const enrichedSchema = Object.assign(schema, {
+    collaborative: true,
+    builtin: "Account",
+  }) as AnyAccountSchema<Shape>;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
 export const coRecordDefiner = <
@@ -203,190 +119,86 @@ export const coRecordDefiner = <
   >;
 };
 
-function enrichCoListSchema<T extends z.core.$ZodType>(schema: z.ZodArray<T>) {
-  const enrichedSchema = Object.assign(schema, {
-    collaborative: true,
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
-      return Object.assign(schema, helpers(schema));
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as CoListSchema<T>;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
 export const coListDefiner = <T extends z.core.$ZodType>(
   element: T,
 ): CoListSchema<T> => {
-  const arraySchema = z.array(element).meta({
+  const schema = z.array(element).meta({
     collaborative: true,
   });
-
-  return enrichCoListSchema(arraySchema);
+  const enrichedSchema = Object.assign(schema, {
+    collaborative: true,
+  }) as AnyCoListSchema<T>;
+  return zodSchemaToCoSchema(enrichedSchema) as unknown as CoListSchema<T>;
 };
 
 export const coProfileDefiner = <
   Shape extends z.core.$ZodLooseShape = Simplify<DefaultProfileShape>,
 >(
-  shape: Shape & {
-    name?: z.core.$ZodString<string>;
-    inbox?: z.core.$ZodOptional<z.core.$ZodString>;
-    inboxInvite?: z.core.$ZodOptional<z.core.$ZodString>;
-  } = {} as any,
+  shape: Shape & Partial<DefaultProfileShape> = {} as any,
 ): CoProfileSchema<Shape> => {
-  const ehnancedShape = Object.assign(shape ?? {}, {
+  const ehnancedShape = Object.assign(shape, {
     name: z.string(),
     inbox: z.optional(z.string()),
     inboxInvite: z.optional(z.string()),
   });
-
   return coMapDefiner(ehnancedShape) as CoProfileSchema<Shape>;
 };
 
-function enrichCoFeedSchema<T extends z.core.$ZodType>(
-  schema: z.ZodCustom<CoFeed<unknown>, unknown>,
+export const coFeedDefiner = <T extends z.core.$ZodType>(
   element: T,
-) {
+): CoFeedSchema<T> => {
+  const schema = z.instanceof(CoFeed);
   const enrichedSchema = Object.assign(schema, {
     collaborative: true,
     builtin: "CoFeed",
     element,
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    withHelpers: (helpers: (Self: z.core.$ZodType) => object) => {
-      return Object.assign(schema, helpers(schema));
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as CoFeedSchema<T>;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
-export const coFeedDefiner = <T extends z.core.$ZodType>(
-  element: T,
-): CoFeedSchema<T> => {
-  return enrichCoFeedSchema(z.instanceof(CoFeed), element);
+  }) as AnyCoFeedSchema<T>;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
-function enrichFileStreamSchema(schema: z.ZodCustom<FileStream, unknown>) {
+export const coFileStreamDefiner = (): FileStreamSchema => {
+  const schema = z.instanceof(FileStream);
   const enrichedSchema = Object.assign(schema, {
     collaborative: true,
     builtin: "FileStream",
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    createFromBlob: (...args: any[]) => {
-      return coSchema.createFromBlob(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    loadAsBlob: (...args: any[]) => {
-      return coSchema.loadAsBlob(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as FileStreamSchema;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
-export const coFileStreamDefiner = (): FileStreamSchema => {
-  return enrichFileStreamSchema(z.instanceof(FileStream));
+  }) as AnyFileStreamSchema;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
-function enrichPlainTextSchema(schema: z.ZodCustom<CoPlainText, unknown>) {
+export const coPlainTextDefiner = (): PlainTextSchema => {
+  const schema = z.instanceof(CoPlainText);
   const enrichedSchema = Object.assign(schema, {
     collaborative: true,
     builtin: "CoPlainText",
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    fromRaw: (...args: any[]) => {
-      return coSchema.fromRaw(...args);
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as PlainTextSchema;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
-export const coPlainTextDefiner = (): PlainTextSchema => {
-  return enrichPlainTextSchema(z.instanceof(CoPlainText));
+  }) as AnyPlainTextSchema;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
-function enrichRichTextSchema(schema: z.ZodCustom<CoRichText, unknown>) {
+export const coRichTextDefiner = (): RichTextSchema => {
+  const schema = z.instanceof(CoRichText);
   const enrichedSchema = Object.assign(schema, {
     collaborative: true,
     builtin: "CoRichText",
-    create: (...args: any[]) => {
-      return coSchema.create(...args);
-    },
-    load: (...args: any[]) => {
-      return coSchema.load(...args);
-    },
-    subscribe: (...args: any[]) => {
-      return coSchema.subscribe(...args);
-    },
-    getCoSchema: () => {
-      return coSchema;
-    },
-  }) as unknown as RichTextSchema;
-
-  // Needs to be derived from the enriched schema
-  const coSchema = zodSchemaToCoSchema(enrichedSchema) as any;
-
-  return enrichedSchema;
-}
-
-export const coRichTextDefiner = (): RichTextSchema => {
-  return enrichRichTextSchema(z.instanceof(CoRichText));
+  }) as AnyRichTextSchema;
+  return zodSchemaToCoSchema(enrichedSchema);
 };
 
 export const coImageDefiner = (): typeof ImageDefinition => {
   return ImageDefinition;
 };
+
+export const coOptionalDefiner = <T extends AnyCoSchema>(
+  schema: T,
+): CoOptionalSchema<T> => {
+  return createCoOptionalSchema(schema);
+};
+
+export const coDiscriminatedUnionDefiner = <
+  T extends readonly [AnyDiscriminableCoSchema, ...AnyDiscriminableCoSchema[]],
+>(
+  discriminator: string,
+  schemas: T,
+): CoDiscriminatedUnionSchema<T> => {
+  const schema = z.discriminatedUnion(discriminator, schemas as any);
+  return zodSchemaToCoSchema(schema);
+};