From 2ce0935f9a7a71b40234f0221caee21da1be3208 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Julien=20Valverd=C3=A9?= <julien.valverde@mailo.com>
Date: Sat, 23 Mar 2024 19:52:44 +0100
Subject: [PATCH] Fixed extend

---
 src/ExtendZodSchemaClass.ts                 | 56 +++++++++++++++++++++
 src/ZodSchemaClass.ts                       |  4 +-
 src/lib.ts                                  |  2 +-
 src/tests.ts                                |  9 ++--
 src/traits/MobXObservableZodSchemaObject.ts |  6 ++-
 5 files changed, 69 insertions(+), 8 deletions(-)
 create mode 100644 src/ExtendZodSchemaClass.ts

diff --git a/src/ExtendZodSchemaClass.ts b/src/ExtendZodSchemaClass.ts
new file mode 100644
index 0000000..9bc2c97
--- /dev/null
+++ b/src/ExtendZodSchemaClass.ts
@@ -0,0 +1,56 @@
+import { TraitInstance, TraitStaticMembers, expression } from "@thilawyn/traitify-ts"
+import { AbstractClass } from "type-fest"
+import { z } from "zod"
+import { ZodSchemaObject, ZodSchemaObjectTrait } from "./lib"
+import { Extend, StaticMembers } from "./util"
+
+
+export function ExtendZodSchemaClass<
+    Self                  extends (
+        AbstractClass<
+            TraitInstance<ZodSchemaObjectTrait<T, Catchall, Values, PartialValues>> & Values,
+            [values: Values]
+        > &
+        TraitStaticMembers<ZodSchemaObjectTrait<T, Catchall, Values, PartialValues>>
+    ),
+
+    T                     extends z.ZodRawShape,
+    Catchall              extends z.ZodTypeAny,
+    Values                extends object,
+    PartialValues         extends Partial<Values>,
+
+    ExtendedT             extends z.ZodRawShape,
+    ExtendedCatchall      extends z.ZodTypeAny,
+    ExtendedValues        extends Values,
+    ExtendedPartialValues extends Partial<ExtendedValues>,
+>(
+    class_: Self | (
+        AbstractClass<
+            TraitInstance<ZodSchemaObjectTrait<T, Catchall, Values, PartialValues>> & Values,
+            [values: Values]
+        > &
+        TraitStaticMembers<ZodSchemaObjectTrait<T, Catchall, Values, PartialValues>>
+    ),
+
+    schemaWithDefaults: (
+        schemaWithDefaults: z.ZodObject<T, "strip", Catchall, Values, PartialValues>
+    ) => z.ZodObject<ExtendedT, "strip", ExtendedCatchall, ExtendedValues, ExtendedPartialValues>,
+) {
+    return expression
+        .extends(
+            class_ as unknown as (
+                AbstractClass<
+                    Omit<
+                        Extend<[InstanceType<Self>, ExtendedValues]>,
+                        keyof TraitInstance<ZodSchemaObjectTrait<T, Catchall, Values, PartialValues>>
+                    >,
+                    [values: ExtendedValues]
+                > &
+                Omit<
+                    StaticMembers<Self>,
+                    keyof TraitStaticMembers<ZodSchemaObjectTrait<T, Catchall, Values, PartialValues>>
+                >
+            )
+        )
+        .expresses(ZodSchemaObject(schemaWithDefaults((class_ as Self).schemaWithDefaults)))
+}
diff --git a/src/ZodSchemaClass.ts b/src/ZodSchemaClass.ts
index 19ee1e9..1b4b9a1 100644
--- a/src/ZodSchemaClass.ts
+++ b/src/ZodSchemaClass.ts
@@ -1,7 +1,7 @@
 import { expression } from "@thilawyn/traitify-ts"
 import { Class } from "type-fest"
 import { z } from "zod"
-import { ExtendableZodSchemaObject } from "./traits/ExtendableZodSchemaObject"
+import { ZodSchemaObject } from "./traits/ZodSchemaObject"
 
 
 export function ZodSchemaClass<
@@ -20,5 +20,5 @@ export function ZodSchemaClass<
                 }
             } as Class<Values, [values: Values]>
         )
-        .expresses(ExtendableZodSchemaObject(schemaWithDefaults))
+        .expresses(ZodSchemaObject(schemaWithDefaults))
 }
diff --git a/src/lib.ts b/src/lib.ts
index 8d3cdc8..b7d4fe2 100644
--- a/src/lib.ts
+++ b/src/lib.ts
@@ -1,6 +1,6 @@
+export { ExtendZodSchemaClass } from "./ExtendZodSchemaClass"
 export { JsonifiedZodSchemaClass } from "./JsonifiedZodSchemaClass"
 export { ZodSchemaClass } from "./ZodSchemaClass"
-export { ExtendableZodSchemaObject } from "./traits/ExtendableZodSchemaObject"
 export { JsonifiedZodSchemaObject } from "./traits/JsonifiedZodSchemaObject"
 export { MobXObservableZodSchemaObject } from "./traits/MobXObservableZodSchemaObject"
 export { ZodSchemaObject, ZodSchemaObjectTrait } from "./traits/ZodSchemaObject"
diff --git a/src/tests.ts b/src/tests.ts
index 57a5ddc..fc6421a 100644
--- a/src/tests.ts
+++ b/src/tests.ts
@@ -1,6 +1,7 @@
 import { Implements } from "@thilawyn/traitify-ts"
 import { Option } from "effect"
 import { z } from "zod"
+import { ExtendZodSchemaClass } from "./ExtendZodSchemaClass"
 import { JsonifiedZodSchemaClass } from "./JsonifiedZodSchemaClass"
 import { ZodSchemaClass } from "./ZodSchemaClass"
 import { s } from "./schema/lib"
@@ -47,9 +48,11 @@ console.log(JSON.stringify(jsonifiedUserInst))
 // jsonifiedUserInst.dejsonify()
 
 
-const adminUserExp = User.extend(s => s.extend({
-    role: z.literal("Admin").default("Admin")
-})).build()
+const adminUserExp = ExtendZodSchemaClass(User,
+    schema => schema.extend({
+        role: z.literal("Admin").default("Admin")
+    })
+).build()
 
 @adminUserExp.staticImplements
 class AdminUser extends adminUserExp.extends implements Implements<typeof adminUserExp> {}
diff --git a/src/traits/MobXObservableZodSchemaObject.ts b/src/traits/MobXObservableZodSchemaObject.ts
index c9913da..9c3a409 100644
--- a/src/traits/MobXObservableZodSchemaObject.ts
+++ b/src/traits/MobXObservableZodSchemaObject.ts
@@ -8,13 +8,15 @@ export const MobXObservableZodSchemaObject = trait
     .staticAbstract(Super => class extends Super {
         declare readonly schema: z.ZodObject<z.ZodRawShape, z.UnknownKeysParam, z.ZodTypeAny, object, object>
     })
-    .implement(Super => class ObservableZodSchemaObject extends Super {
+    .implement(Super => class MobXObservableZodSchemaObjectImpl extends Super {
+        declare ["constructor"]: typeof MobXObservableZodSchemaObjectImpl
+
         constructor(...args: any[]) {
             super(...args)
 
             makeObservable(this,
                 mapValues(
-                    (this.constructor as typeof ObservableZodSchemaObject).schema.shape,
+                    this.constructor.schema.shape,
                     () => observable,
                 )
             )