Refactor OnyxUpdate type to bring back type safety

.eslintignore

@@ -1,4 +1,6 @@
 *.d.ts
 dist
 node_modules
-*.config.js
+*.config.js
+# tests/types catalog is not type checked with the rest of the project, so we need to ignore it in eslint
+tests/types/**/*.ts

.github/workflows/test.yml

@@ -30,3 +30,5 @@ jobs:
             - run: npm run test
               env:
                   CI: true
+
+            - run: npm run test:types

jest.config.js

@@ -4,7 +4,7 @@ module.exports = {
     transform: {
         '\\.[jt]sx?$': 'babel-jest',
     },
-    testPathIgnorePatterns: ['<rootDir>/node_modules/', '<rootDir>/tests/unit/mocks/', '<rootDir>/tests/e2e/'],
+    testPathIgnorePatterns: ['<rootDir>/node_modules/', '<rootDir>/tests/unit/mocks/', '<rootDir>/tests/e2e/', '<rootDir>/tests/types/'],
     testMatch: ['**/tests/unit/**/*.[jt]s?(x)', '**/?(*.)+(spec|test).[jt]s?(x)'],
     globals: {
         __DEV__: true,

lib/Onyx.ts

@@ -4,8 +4,6 @@ import Storage from './storage';
 import utils from './utils';
 import DevTools, {initDevTools} from './DevTools';
 import type {
-    Collection,
-    CollectionKey,
     CollectionKeyBase,
     ConnectOptions,
     InitOptions,
@@ -15,6 +13,7 @@ import type {
     MixedOperationsQueue,
     OnyxKey,
     OnyxMergeCollectionInput,
+    OnyxSetCollectionInput,
     OnyxMergeInput,
     OnyxMultiSetInput,
     OnyxSetInput,
@@ -374,7 +373,7 @@ function merge<TKey extends OnyxKey>(key: TKey, changes: OnyxMergeInput<TKey>):
  * @param collectionKey e.g. `ONYXKEYS.COLLECTION.REPORT`
  * @param collection Object collection keyed by individual collection member keys and values
  */
-function mergeCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: OnyxMergeCollectionInput<TKey, TMap>): Promise<void> {
+function mergeCollection<TKey extends CollectionKeyBase>(collectionKey: TKey, collection: OnyxMergeCollectionInput<TKey>): Promise<void> {
     return OnyxUtils.mergeCollectionWithPatches(collectionKey, collection, undefined, true);
 }
 
@@ -545,7 +544,7 @@ function update(data: OnyxUpdate[]): Promise<void> {
             [OnyxUtils.METHOD.SET]: enqueueSetOperation,
             [OnyxUtils.METHOD.MERGE]: enqueueMergeOperation,
             [OnyxUtils.METHOD.MERGE_COLLECTION]: () => {
-                const collection = value as Collection<CollectionKey, unknown, unknown>;
+                const collection = value as OnyxMergeCollectionInput<OnyxKey>;
                 if (!OnyxUtils.isValidNonEmptyCollectionForMerge(collection)) {
                     Logger.logInfo('mergeCollection enqueued within update() with invalid or empty value. Skipping this operation.');
                     return;
@@ -558,7 +557,7 @@ function update(data: OnyxUpdate[]): Promise<void> {
                     collectionKeys.forEach((collectionKey) => enqueueMergeOperation(collectionKey, mergedCollection[collectionKey]));
                 }
             },
-            [OnyxUtils.METHOD.SET_COLLECTION]: (k, v) => promises.push(() => setCollection(k, v as Collection<CollectionKey, unknown, unknown>)),
+            [OnyxUtils.METHOD.SET_COLLECTION]: (k, v) => promises.push(() => setCollection(k, v as OnyxSetCollectionInput<OnyxKey>)),
             [OnyxUtils.METHOD.MULTI_SET]: (k, v) => Object.entries(v as Partial<OnyxInputKeyValueMapping>).forEach(([entryKey, entryValue]) => enqueueSetOperation(entryKey, entryValue)),
             [OnyxUtils.METHOD.CLEAR]: () => {
                 clearPromise = clear();
@@ -611,14 +610,14 @@ function update(data: OnyxUpdate[]): Promise<void> {
             promises.push(() =>
                 OnyxUtils.mergeCollectionWithPatches(
                     collectionKey,
-                    batchedCollectionUpdates.merge as Collection<CollectionKey, unknown, unknown>,
+                    batchedCollectionUpdates.merge as OnyxMergeCollectionInput<OnyxKey>,
                     batchedCollectionUpdates.mergeReplaceNullPatches,
                     true,
                 ),
             );
         }
         if (!utils.isEmptyObject(batchedCollectionUpdates.set)) {
-            promises.push(() => OnyxUtils.partialSetCollection(collectionKey, batchedCollectionUpdates.set as Collection<CollectionKey, unknown, unknown>));
+            promises.push(() => OnyxUtils.partialSetCollection(collectionKey, batchedCollectionUpdates.set as OnyxSetCollectionInput<OnyxKey>));
         }
     });
 
@@ -655,7 +654,7 @@ function update(data: OnyxUpdate[]): Promise<void> {
  * @param collectionKey e.g. `ONYXKEYS.COLLECTION.REPORT`
  * @param collection Object collection keyed by individual collection member keys and values
  */
-function setCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: OnyxMergeCollectionInput<TKey, TMap>): Promise<void> {
+function setCollection<TKey extends CollectionKeyBase>(collectionKey: TKey, collection: OnyxSetCollectionInput<TKey>): Promise<void> {
     let resultCollection: OnyxInputKeyValueMapping = collection;
     let resultCollectionKeys = Object.keys(resultCollection);
 

lib/OnyxUtils.ts

@@ -29,6 +29,7 @@ import type {
     OnyxUpdate,
     OnyxValue,
     Selector,
+    OnyxSetCollectionInput,
 } from './types';
 import type {FastMergeOptions, FastMergeResult} from './utils';
 import utils from './utils';
@@ -1035,7 +1036,7 @@ function initializeWithDefaultKeyStates(): Promise<void> {
 /**
  * Validate the collection is not empty and has a correct type before applying mergeCollection()
  */
-function isValidNonEmptyCollectionForMerge<TKey extends CollectionKeyBase, TMap>(collection: OnyxMergeCollectionInput<TKey, TMap>): boolean {
+function isValidNonEmptyCollectionForMerge<TKey extends CollectionKeyBase>(collection: OnyxMergeCollectionInput<TKey>): boolean {
     return typeof collection === 'object' && !Array.isArray(collection) && !utils.isEmptyObject(collection);
 }
 
@@ -1241,9 +1242,9 @@ function updateSnapshots(data: OnyxUpdate[], mergeFn: typeof Onyx.merge): Array<
  * @param mergeReplaceNullPatches Record where the key is a collection member key and the value is a list of
  * tuples that we'll use to replace the nested objects of that collection member record with something else.
  */
-function mergeCollectionWithPatches<TKey extends CollectionKeyBase, TMap>(
+function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
     collectionKey: TKey,
-    collection: OnyxMergeCollectionInput<TKey, TMap>,
+    collection: OnyxMergeCollectionInput<TKey>,
     mergeReplaceNullPatches?: MultiMergeReplaceNullPatches,
     isProcessingCollectionUpdate = false,
 ): Promise<void> {
@@ -1350,7 +1351,7 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase, TMap>(
             });
 
             return Promise.all(promises)
-                .catch((error) => evictStorageAndRetry(error, mergeCollectionWithPatches, collectionKey, resultCollection))
+                .catch((error) => evictStorageAndRetry(error, mergeCollectionWithPatches, collectionKey, resultCollection as OnyxMergeCollectionInput<TKey>))
                 .then(() => {
                     sendActionToDevTools(METHOD.MERGE_COLLECTION, undefined, resultCollection);
                     return promiseUpdate;
@@ -1366,7 +1367,7 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase, TMap>(
  * @param collectionKey e.g. `ONYXKEYS.COLLECTION.REPORT`
  * @param collection Object collection keyed by individual collection member keys and values
  */
-function partialSetCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: OnyxMergeCollectionInput<TKey, TMap>): Promise<void> {
+function partialSetCollection<TKey extends CollectionKeyBase>(collectionKey: TKey, collection: OnyxSetCollectionInput<TKey>): Promise<void> {
     let resultCollection: OnyxInputKeyValueMapping = collection;
     let resultCollectionKeys = Object.keys(resultCollection);
 

lib/index.ts

@@ -16,6 +16,7 @@ import type {
     OnyxMultiSetInput,
     OnyxMergeInput,
     OnyxMergeCollectionInput,
+    OnyxSetCollectionInput,
 } from './types';
 import type {FetchStatus, ResultMetadata, UseOnyxResult, UseOnyxOptions} from './useOnyx';
 import type {Connection} from './OnyxConnectionManager';
@@ -40,6 +41,7 @@ export type {
     OnyxMultiSetInput,
     OnyxMergeInput,
     OnyxMergeCollectionInput,
+    OnyxSetCollectionInput,
     OnyxUpdate,
     OnyxValue,
     ResultMetadata,

lib/types.ts

@@ -1,5 +1,4 @@
 import type {Merge} from 'type-fest';
-import type {BuiltIns} from 'type-fest/source/internal';
 import type OnyxUtils from './OnyxUtils';
 import type {OnyxMethod} from './OnyxUtils';
 import type {FastMergeReplaceNullPatch} from './utils';
@@ -157,6 +156,10 @@ type OnyxValue<TKey extends OnyxKey> = string extends TKey ? unknown : TKey exte
 /** Utility type to extract `TOnyxValue` from `OnyxCollection<TOnyxValue>` */
 type ExtractOnyxCollectionValue<TOnyxCollection> = TOnyxCollection extends NonNullable<OnyxCollection<infer U>> ? U : never;
 
+type Primitive = null | undefined | string | number | boolean | symbol | bigint;
+
+type BuiltIns = Primitive | void | Date | RegExp;
+
 type NonTransformableTypes =
     | BuiltIns
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -205,13 +208,7 @@ type NullishObjectDeep<ObjectType extends object> = {
  * Also, the `TMap` type is inferred automatically in `mergeCollection()` method and represents
  * the object of collection keys/values specified in the second parameter of the method.
  */
-type Collection<TKey extends CollectionKeyBase, TValue, TMap = never> = {
-    [MapK in keyof TMap]: MapK extends `${TKey}${string}`
-        ? MapK extends `${TKey}`
-            ? never // forbids empty id
-            : TValue
-        : never;
-};
+type Collection<TKey extends CollectionKeyBase, TValue> = Record<`${TKey}${string}`, TValue> & {[P in TKey]?: never};
 
 /** Represents the base options used in `Onyx.connect()` method. */
 // NOTE: Any changes to this type like adding or removing options must be accounted in OnyxConnectionManager's `generateConnectionID()` method!
@@ -322,48 +319,58 @@ type OnyxMergeInput<TKey extends OnyxKey> = OnyxInput<TKey>;
 /**
  * This represents the value that can be passed to `Onyx.merge` and to `Onyx.update` with the method "MERGE"
  */
-type OnyxMergeCollectionInput<TKey extends OnyxKey, TMap = object> = Collection<TKey, NonNullable<OnyxInput<TKey>>, TMap>;
+type OnyxMergeCollectionInput<TKey extends OnyxKey> = Collection<TKey, NonNullable<OnyxInput<TKey>>>;
 
-type OnyxMethodMap = typeof OnyxUtils.METHOD;
+/**
+ * This represents the value that can be passed to `Onyx.setCollection` and to `Onyx.update` with the method "SET_COLLECTION"
+ */
+type OnyxSetCollectionInput<TKey extends OnyxKey> = Collection<TKey, OnyxInput<TKey>>;
 
-// Maps onyx methods to their corresponding value types
-type OnyxMethodValueMap = {
-    [OnyxUtils.METHOD.SET]: {
-        key: OnyxKey;
-        value: OnyxSetInput<OnyxKey>;
-    };
-    [OnyxUtils.METHOD.MULTI_SET]: {
-        key: OnyxKey;
-        value: OnyxMultiSetInput;
-    };
-    [OnyxUtils.METHOD.MERGE]: {
-        key: OnyxKey;
-        value: OnyxMergeInput<OnyxKey>;
-    };
-    [OnyxUtils.METHOD.CLEAR]: {
-        key: OnyxKey;
-        value?: undefined;
-    };
-    [OnyxUtils.METHOD.MERGE_COLLECTION]: {
-        key: CollectionKeyBase;
-        value: OnyxMergeCollectionInput<CollectionKeyBase>;
-    };
-    [OnyxUtils.METHOD.SET_COLLECTION]: {
-        key: CollectionKeyBase;
-        value: OnyxMergeCollectionInput<CollectionKeyBase>;
-    };
-};
+type OnyxMethodMap = typeof OnyxUtils.METHOD;
 
 /**
  * OnyxUpdate type includes all onyx methods used in OnyxMethodValueMap.
  * If a new method is added to OnyxUtils.METHOD constant, it must be added to OnyxMethodValueMap type.
  * Otherwise it will show static type errors.
  */
-type OnyxUpdate = {
-    [Method in OnyxMethod]: {
-        onyxMethod: Method;
-    } & OnyxMethodValueMap[Method];
-}[OnyxMethod];
+type OnyxUpdate =
+    // ⚠️ DO NOT CHANGE THIS TYPE, UNLESS YOU KNOW WHAT YOU ARE DOING. ⚠️
+    | {
+          [TKey in OnyxKey]:
+              | {
+                    onyxMethod: typeof OnyxUtils.METHOD.SET;
+                    key: TKey;
+                    value: OnyxSetInput<TKey>;
+                }
+              | {
+                    onyxMethod: typeof OnyxUtils.METHOD.MULTI_SET;
+                    key: TKey;
+                    value: OnyxMultiSetInput;
+                }
+              | {
+                    onyxMethod: typeof OnyxUtils.METHOD.MERGE;
+                    key: TKey;
+                    value: OnyxMergeInput<TKey>;
+                }
+              | {
+                    onyxMethod: typeof OnyxUtils.METHOD.CLEAR;
+                    key: TKey;
+                    value?: undefined;
+                };
+      }[OnyxKey]
+    | {
+          [TKey in CollectionKeyBase]:
+              | {
+                    onyxMethod: typeof OnyxUtils.METHOD.MERGE_COLLECTION;
+                    key: TKey;
+                    value: OnyxMergeCollectionInput<TKey>;
+                }
+              | {
+                    onyxMethod: typeof OnyxUtils.METHOD.SET_COLLECTION;
+                    key: TKey;
+                    value: OnyxSetCollectionInput<TKey>;
+                };
+      }[CollectionKeyBase];
 
 /**
  * Represents the options used in `Onyx.set()` method.
@@ -474,6 +481,7 @@ export type {
     OnyxMultiSetInput,
     OnyxMergeInput,
     OnyxMergeCollectionInput,
+    OnyxSetCollectionInput,
     OnyxMethod,
     OnyxMethodMap,
     OnyxUpdate,

package.json

@@ -32,6 +32,7 @@
     "lint": "eslint .",
     "typecheck": "tsc --noEmit",
     "test": "jest",
+    "test:types": "npm run build && tsc --noEmit --project tsconfig.test.json",
     "perf-test": "npx reassure",
     "build": "tsc -p tsconfig.build.json",
     "build:watch": "nodemon --watch lib --ext js,json,ts,tsx --exec \"npm run build && npm pack\"",

tests/types/OnyxUpdate.ts

@@ -0,0 +1,59 @@
+import type {OnyxUpdate} from '../../dist/types';
+import ONYX_KEYS from './setup';
+
+const onyxUpdate: OnyxUpdate = {
+    onyxMethod: 'set',
+    key: ONYX_KEYS.TEST_KEY,
+    value: 'string',
+};
+
+const onyxUpdateError: OnyxUpdate = {
+    onyxMethod: 'set',
+    key: ONYX_KEYS.TEST_KEY,
+    // @ts-expect-error TEST_KEY is a string, not a number
+    value: 2,
+};
+
+const onyxUpdateCollection: OnyxUpdate = {
+    onyxMethod: 'mergecollection',
+    key: ONYX_KEYS.COLLECTION.TEST_KEY,
+    value: {
+        [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: {
+            str: 'test',
+        },
+        [`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: {
+            str: 'test2',
+        },
+    },
+};
+
+// @ts-expect-error COLLECTION.TEST_KEY is an object, not a number
+const onyxUpdateCollectionError: OnyxUpdate = {
+    onyxMethod: 'mergecollection',
+    key: ONYX_KEYS.COLLECTION.TEST_KEY,
+    value: {
+        [`${ONYX_KEYS.COLLECTION.TEST_KEY}1`]: 2,
+    },
+};
+
+const onyxUpdateCollectionError2: OnyxUpdate = {
+    onyxMethod: 'mergecollection',
+    key: ONYX_KEYS.COLLECTION.TEST_KEY,
+    value: {
+        [`${ONYX_KEYS.COLLECTION.TEST_KEY}2`]: {
+            // @ts-expect-error nonExistingKey is not a valid key
+            nonExistingKey: 'test2',
+        },
+    },
+};
+
+// @ts-expect-error COLLECTION.TEST_KEY is invalid key, it is missing the suffix
+const onyxUpdateCollectionError3: OnyxUpdate = {
+    onyxMethod: 'mergecollection',
+    key: ONYX_KEYS.COLLECTION.TEST_KEY,
+    value: {
+        [ONYX_KEYS.COLLECTION.TEST_KEY]: {
+            str: 'test2',
+        },
+    },
+};

tests/types/mergeCollection.ts

@@ -0,0 +1,26 @@
+import Onyx from '../../dist/Onyx';
+import ONYX_KEYS from './setup';
+
+Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
+    test_1: {
+        str: 'test3',
+    },
+    test_2: {
+        str: 'test4',
+    },
+    test_3: {
+        str: 'test5',
+    },
+});
+
+Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_KEY, {
+    // @ts-expect-error COLLECTION.TEST_KEY is invalid key, it is missing the suffix
+    test_: {
+        str: 'test3',
+    },
+    test_2: {
+        str: 'test4',
+    },
+    // @ts-expect-error COLLECTION.TEST_KEY is object, not a number
+    test_3: 2,
+});