feat: support lookup table with number key

LegendApp · Jul 1, 2024 · 2ff99ae · 2ff99ae
1 parent 730ff1c
commit 2ff99ae
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 8 deletions.
diff --git a/src/observableTypes.ts b/src/observableTypes.ts
@@ -36,8 +36,8 @@ export type RemoveObservables<T> =
               ? TObs
               : T extends () => infer TRet
                 ? RemoveObservables<TRet> & T
-                : T extends (key: string) => infer TRet
-                  ? Record<string, RemoveObservables<TRet>> & T
+                : T extends (key: infer TKey extends string | number) => infer TRet
+                  ? Record<TKey, RemoveObservables<TRet>> & T
                   : T;
 
 interface ObservableArray<T, U>
@@ -137,8 +137,8 @@ type ObservableChildren<T, Nullable = IsNullable<T>> = {
 type ObservableFunctionChildren<T> = {
     [K in keyof T]-?: T[K] extends Observable
         ? T[K]
-        : T[K] extends (key: infer Key extends string) => Promise<infer t> | infer t
-          ? HasOneStringParam<T[K]> extends true
+        : T[K] extends (key: infer Key extends string | number) => Promise<infer t> | infer t
+          ? IsLookupFunction<T[K]> extends true
               ? Observable<Record<Key, t>> & T[K]
               : t extends void
                 ? T[K]
@@ -157,9 +157,9 @@ export type ObservableObject<T> = ObservableObjectFunctions<ObservableProps<T> &
 type ObservableFunction<T> = T extends () => infer t ? t | (() => t) : T;
 
 // Check if the function type T has one lookup parameter
-type HasOneStringParam<T> = T extends (...args: infer P) => any
+type IsLookupFunction<T> = T extends (...args: infer P) => any
     ? P extends { length: 1 }
-        ? P[0] extends string | ObservablePrimitive<string>
+        ? P[0] extends string | ObservablePrimitive<string> | number | ObservablePrimitive<number>
             ? true
             : false
         : false
@@ -175,10 +175,10 @@ type ObservableNode<T, NT = NonNullable<T>> = [NT] extends [never] // means that
         ? ObservableNode<t>
         : [T] extends [(key: infer K extends string) => infer t]
           ? [t] extends [ImmutableObservableBase<any>]
-              ? HasOneStringParam<T> extends true
+              ? IsLookupFunction<T> extends true
                   ? Observable<Record<K, t>>
                   : t
-              : HasOneStringParam<T> extends true
+              : IsLookupFunction<T> extends true
                 ? Observable<Record<K, t>> & T
                 : Observable<ObservableFunction<t>>
           : [NT] extends [ImmutableObservableBase<any>]

diff --git a/tests/types.test.ts b/tests/types.test.ts
@@ -573,12 +573,37 @@ describe('lookup table', () => {
         });
 
         const gotten = obs$.get();
+        const child = obs$.lookup['hi'].get();
 
         expectTypeOf<(typeof obs$)['computed']['get']>().returns.toMatchTypeOf<Record<string, string>>();
         expectTypeOf<(typeof obs$)['lookup']['get']>().returns.toMatchTypeOf<Record<string, string>>();
         expectTypeOf<(typeof gotten)['computed']>().toMatchTypeOf<Record<string, string>>();
         expectTypeOf<(typeof gotten)['computed']>().toMatchTypeOf<() => ObservablePrimitive<Record<string, string>>>();
         expectTypeOf<(typeof gotten)['lookup']>().toMatchTypeOf<Record<string, string>>();
         expectTypeOf<(typeof gotten)['lookup']>().toMatchTypeOf<(key: string) => Observable<string>>();
+        expectTypeOf<typeof child>().toEqualTypeOf<string>();
+    });
+    it('should type lookup table as Record', () => {
+        const num$ = observable<Record<number, string>>({ 0: 'hi' });
+
+        const obs$ = observable({
+            computed: () => {
+                return num$;
+            },
+            lookup: (x: number) => {
+                return num$[x];
+            },
+        });
+
+        const gotten = obs$.get();
+        const child = obs$.lookup[0].get();
+
+        expectTypeOf<(typeof obs$)['computed']['get']>().returns.toMatchTypeOf<Record<number, string>>();
+        expectTypeOf<(typeof obs$)['lookup']['get']>().returns.toMatchTypeOf<Record<number, string>>();
+        expectTypeOf<(typeof gotten)['computed']>().toMatchTypeOf<Record<number, string>>();
+        expectTypeOf<(typeof gotten)['computed']>().toMatchTypeOf<() => ObservablePrimitive<Record<number, string>>>();
+        expectTypeOf<(typeof gotten)['lookup']>().toMatchTypeOf<Record<number, string>>();
+        expectTypeOf<(typeof gotten)['lookup']>().toMatchTypeOf<(key: number) => Observable<string>>();
+        expectTypeOf<typeof child>().toEqualTypeOf<string>();
     });
 });