From e7c65c0a07b193b87beffe9bdcf4182cc44aa6f7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Julien=20Valverd=C3=A9?= <julien.valverde@mailo.com>
Date: Sun, 25 Feb 2024 04:32:09 +0100
Subject: [PATCH] ZodSchemaClassExtender

---
 src/builders/ZodSchemaClassBuilder.ts  | 18 ++++----
 src/builders/ZodSchemaClassExtender.ts | 59 ++++++++++++++++++++++++++
 src/util/extend.ts                     |  8 ++--
 3 files changed, 72 insertions(+), 13 deletions(-)
 create mode 100644 src/builders/ZodSchemaClassExtender.ts

diff --git a/src/builders/ZodSchemaClassBuilder.ts b/src/builders/ZodSchemaClassBuilder.ts
index 75b85f1..fe12b66 100644
--- a/src/builders/ZodSchemaClassBuilder.ts
+++ b/src/builders/ZodSchemaClassBuilder.ts
@@ -18,7 +18,7 @@ extends TraitExpressionBuilder<Superclass, Traits> {
     declare ["constructor"]: typeof ZodSchemaClassBuilder
 
     schema<
-        Superclass        extends AbstractClass<object, []> & { schema?: never, defaultValues?: never },
+        Super             extends AbstractClass<object, []> & { schema?: never, defaultValues?: never },
 
         SchemaT           extends z.ZodRawShape,
         SchemaUnknownKeys extends z.UnknownKeysParam,
@@ -27,7 +27,7 @@ extends TraitExpressionBuilder<Superclass, Traits> {
         Values            extends object,
         DefaultValues     extends Partial<Values>,
     >(
-        this: ZodSchemaClassBuilder<Superclass, Traits>,
+        this: ZodSchemaClassBuilder<Super, Traits>,
 
         { schema, defaultValues }: {
             schema:         z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>
@@ -47,10 +47,10 @@ extends TraitExpressionBuilder<Superclass, Traits> {
         return new this.constructor(
             Schemas as unknown as (
                 AbstractClass<
-                    InstanceType<Superclass> & Values,
+                    InstanceType<Super> & Values,
                     ConstructorParameters<typeof Schemas>
                 > &
-                StaticMembers<Superclass> &
+                StaticMembers<Super> &
                 StaticMembers<typeof Schemas>
             ),
 
@@ -63,7 +63,7 @@ extends TraitExpressionBuilder<Superclass, Traits> {
     }
 
     jsonifiable<
-        Superclass                 extends ZodSchemaAbstractClass<Instance, SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, DefaultValues>
+        Super                      extends ZodSchemaAbstractClass<Instance, SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, DefaultValues>
                                          & { jsonifySchema?: never, dejsonifySchema?: never },
         Instance                   extends Values,
 
@@ -85,7 +85,7 @@ extends TraitExpressionBuilder<Superclass, Traits> {
         JsonifiedValues            extends JsonifiableObject,
     >(
         this: ZodSchemaClassBuilder<
-            Superclass | ZodSchemaAbstractClass<Instance, SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, DefaultValues>,
+            Super | ZodSchemaAbstractClass<Instance, SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, DefaultValues>,
             Traits
         >,
 
@@ -112,12 +112,12 @@ extends TraitExpressionBuilder<Superclass, Traits> {
         return new this.constructor(
             JsonifiableSchemas as unknown as (
                 AbstractClass<
-                    InstanceType<Superclass> & JsonifiableSchemas,
+                    Instance & JsonifiableSchemas,
 
-                    // TODO: for some reason, ConstructorParameters<Superclass> does not work here. Maybe try to find a fix?
+                    // TODO: for some reason, ConstructorParameters<Super> does not work here. Maybe try to find a fix?
                     ConstructorParameters<ZodSchemaAbstractClass<Instance, SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, DefaultValues>>
                 > &
-                StaticMembers<Superclass> &
+                StaticMembers<Super> &
                 StaticMembers<typeof JsonifiableSchemas>
             ),
 
diff --git a/src/builders/ZodSchemaClassExtender.ts b/src/builders/ZodSchemaClassExtender.ts
new file mode 100644
index 0000000..d365476
--- /dev/null
+++ b/src/builders/ZodSchemaClassExtender.ts
@@ -0,0 +1,59 @@
+import { AbstractClass } from "type-fest"
+import { z } from "zod"
+import { ZodSchemaAbstractClass } from "../shapes/ZodSchemaClass"
+import { Extend, StaticMembers } from "../util"
+
+
+export class ZodSchemaClassExtender<Superclass extends AbstractClass<object>> {
+    declare ["constructor"]: typeof ZodSchemaClassExtender
+
+    constructor(readonly superclass: Superclass) {}
+
+
+    schema<
+        Super                  extends ZodSchemaAbstractClass<SuperInstance, SuperSchemaT, SuperSchemaUnknownKeys, SuperSchemaCatchall, SuperValues, SuperDefaultValues>,
+        SuperInstance          extends SuperValues,
+
+        SuperSchemaT           extends z.ZodRawShape,
+        SuperSchemaUnknownKeys extends z.UnknownKeysParam,
+        SuperSchemaCatchall    extends z.ZodTypeAny,
+
+        SuperValues            extends object,
+        SuperDefaultValues     extends Partial<SuperValues>,
+
+        SchemaT                extends z.ZodRawShape,
+        SchemaUnknownKeys      extends z.UnknownKeysParam,
+        SchemaCatchall         extends z.ZodTypeAny,
+
+        Values                 extends SuperValues,
+        DefaultValues          extends Partial<Values>,
+    >(
+        this: ZodSchemaClassExtender<
+            Super | ZodSchemaAbstractClass<SuperInstance, SuperSchemaT, SuperSchemaUnknownKeys, SuperSchemaCatchall, SuperValues, SuperDefaultValues>
+        >,
+
+        props: {
+            schema:         (schema: Super["schema"]) => z.ZodObject<SchemaT, SchemaUnknownKeys, SchemaCatchall, Values, Values>
+            defaultValues:  (defaultValues: SuperDefaultValues) => DefaultValues
+        },
+    ) {
+        const schema        = props.schema(this.superclass.schema)
+        const defaultValues = props.defaultValues(this.superclass.defaultValues)
+
+        class Schemas extends (this.superclass as AbstractClass<object, any[]>) {
+            static readonly schema        = schema
+            static readonly defaultValues = defaultValues
+        }
+
+        return new this.constructor(
+            this.superclass as unknown as AbstractClass<
+                Extend<[SuperInstance, Schemas]>,
+                [values: Values]
+            > &
+            Extend<[
+                StaticMembers<Super>,
+                StaticMembers<typeof Schemas>,
+            ]>
+        )
+    }
+}
diff --git a/src/util/extend.ts b/src/util/extend.ts
index baa104d..5cab210 100644
--- a/src/util/extend.ts
+++ b/src/util/extend.ts
@@ -6,18 +6,18 @@
 export type CommonKeys<A, B> = Extract<keyof A, keyof B>
 
 export type Extend<T extends readonly object[]> = (
-    T extends [
+    T extends readonly [
         infer Super,
         infer Self,
-        ...infer Rest extends object[],
+        ...infer Rest extends readonly object[],
     ]
         ? Pick<Self, CommonKeys<Self, Super>> extends Pick<Super, CommonKeys<Self, Super>>
-            ? Extend<[
+            ? Extend<readonly [
                 Omit<Super, CommonKeys<Self, Super>> & Self,
                 ...Rest,
             ]>
             : never
-        : T extends [infer Self]
+        : T extends readonly [infer Self]
             ? Self
             : {}
 )