From 01615aff2129bc4fe9d58f5c597c6046c2d2f22f Mon Sep 17 00:00:00 2001
From: Endel Dreyer <endel.dreyer@gmail.com>
Date: Fri, 18 Aug 2023 18:53:07 -0300
Subject: [PATCH] improve toJSON() return type.

---
 package.json             |  2 +-
 src/Schema.ts            |  6 +++---
 src/types/HelperTypes.ts | 40 ++++++++++++----------------------------
 test/SchemaTest.ts       |  2 +-
 4 files changed, 17 insertions(+), 33 deletions(-)

diff --git a/package.json b/package.json
index ce311ff3..5bbf5441 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@colyseus/schema",
-  "version": "2.0.12",
+  "version": "2.0.13",
   "description": "Binary state serializer with delta encoding for games",
   "bin": {
     "schema-codegen": "./bin/schema-codegen"
diff --git a/src/Schema.ts b/src/Schema.ts
index de63cdd2..25c6bd5b 100644
--- a/src/Schema.ts
+++ b/src/Schema.ts
@@ -11,7 +11,7 @@ import { CollectionSchema } from './types/CollectionSchema';
 import { SetSchema } from './types/SetSchema';
 
 import { ChangeTree, Ref, ChangeOperation } from "./changes/ChangeTree";
-import { NonFunctionPropNames } from './types/HelperTypes';
+import { NonFunctionPropNames, ToJSON } from './types/HelperTypes';
 import { ClientState } from './filters';
 import { getType } from './types/typeRegistry';
 import { ReferenceTracker } from './changes/ReferenceTracker';
@@ -184,7 +184,7 @@ export abstract class Schema {
     }
 
     public assign(
-        props: { [prop in NonFunctionPropNames<this>]?: this[prop] }
+        props: { [prop in NonFunctionPropNames<this>]?: this[prop] } | ToJSON<this>,
     ) {
         Object.assign(this, props);
         return this;
@@ -898,7 +898,7 @@ export abstract class Schema {
         const schema = this._definition.schema;
         const deprecated = this._definition.deprecated;
 
-        const obj = {}
+        const obj: ToJSON<typeof this> = {};
         for (let field in schema) {
             if (!deprecated[field] && this[field] !== null && typeof (this[field]) !== "undefined") {
                 obj[field] = (typeof (this[field]['toJSON']) === "function")
diff --git a/src/types/HelperTypes.ts b/src/types/HelperTypes.ts
index f4d47a6d..a221850d 100644
--- a/src/types/HelperTypes.ts
+++ b/src/types/HelperTypes.ts
@@ -1,34 +1,18 @@
-type Bool = 'true' | 'false'
-type Key = string | number | symbol;
+import { ArraySchema } from "./ArraySchema";
+import { MapSchema } from "./MapSchema";
 
-type Not<X extends Bool> = {
-    true: 'false',
-    false: 'true'
-}[X]
-
-type HaveIntersection<S1 extends string, S2 extends string> = (
-    { [K in S1]: 'true' } &
-    { [key: string]: 'false' }
-)[S2]
-
-type IsNeverWorker<S extends Key> = (
-    { [K in S]: 'false' } &
-    { [key: string]: 'true' }
-)[S]
-
-// Worker needed because of https://github.com/Microsoft/TypeScript/issues/18118
-type IsNever<T extends Key> = Not<HaveIntersection<IsNeverWorker<T>, 'false'>>
-
-type IsFunction<T> = IsNever<keyof T>
-
-export type NonFunctionProps<T> = {
-    [K in keyof T]: {
-        'false': K,
-        'true': never
-    }[IsFunction<T[K]>]
-}[keyof T];
+export type NonFunctionProps<T> = Omit<T, {
+    [K in keyof T]: T[K] extends Function ? K : never;
+}[keyof T]>;
 
 export type NonFunctionPropNames<T> = {
     [K in keyof T]: T[K] extends Function ? never : K
 }[keyof T];
 
+export type ToJSON<T> = Partial<NonFunctionProps<{
+    [K in keyof T]: T[K] extends MapSchema<infer U>
+        ? Record<string, U>
+        : T[K] extends ArraySchema<infer U>
+            ? U[]
+            : T[K]
+}>>;
\ No newline at end of file
diff --git a/test/SchemaTest.ts b/test/SchemaTest.ts
index 92394051..20492d2d 100644
--- a/test/SchemaTest.ts
+++ b/test/SchemaTest.ts
@@ -973,7 +973,7 @@ describe("Schema Usage", () => {
 
             state.value.set('k1', 1);
             const firstEncoded = state.encodeAll();
-            
+
             const decodedState1 = new TestMapSchema();
             decodedState1.decode(firstEncoded);
             assert.deepStrictEqual(decodedState1.toJSON(), state.toJSON(), `decodedState1.toJSON() EQ state.toJSON()`);