fix(atoms): make InjectPromiseAtomApi methods return the subtype (#320)

Author: bowheart

packages/atoms/src/classes/AtomApi.ts

@@ -125,9 +125,9 @@ export class AtomApi<G extends AtomApiGenerics> {
     return this as unknown as AtomApi<Omit<G, 'Promise'> & { Promise: P }> // for chaining
   }
 
-  public setTtl(ttl: AtomInstanceTtl | (() => AtomInstanceTtl)) {
+  public setTtl(ttl: AtomInstanceTtl | (() => AtomInstanceTtl)): AtomApi<G> {
     this.ttl = ttl
 
-    return this // for chaining
+    return this as AtomApi<G> // for chaining
   }
 }

packages/atoms/src/factories/ion.ts

@@ -107,34 +107,6 @@ export const ion: {
     ResolvedState: ResolvedState
   }>
 
-  // No Signal (TODO: Is this overload unnecessary? `atom` doesn't have it)
-  <
-    State = any,
-    Params extends any[] = [],
-    Exports extends Record<string, any> = None,
-    PromiseType extends AtomApiPromise = undefined
-  >(
-    key: string,
-    value: (
-      ecosystem: Ecosystem,
-      ...params: Params
-    ) =>
-      | AtomApi<{
-          Exports: Exports
-          Promise: PromiseType
-          Signal: undefined
-          State: State
-        }>
-      | State,
-    config?: AtomConfig<State>
-  ): IonTemplateRecursive<{
-    State: State
-    Params: Params
-    Events: None
-    Exports: Exports
-    Promise: PromiseType
-  }>
-
   // Catch-all
   <
     State = any,

packages/atoms/src/injectors/injectPromise.ts

@@ -6,11 +6,14 @@ import {
 } from '../utils/promiseUtils'
 import {
   AtomApiGenerics,
+  AtomInstanceTtl,
+  ExportsConfig,
   InjectorDeps,
   InjectPromiseConfig,
   InjectSignalConfig,
   InternalEvaluationReason,
   None,
+  Prettify,
   PromiseState,
 } from '../types/index'
 import { injectEffect } from './injectEffect'
@@ -34,6 +37,43 @@ export interface InjectPromiseAtomApi<
     ResolvedState: Data
     State: Data | undefined
   }>
+
+  addExports<NewExports extends Record<string, any>>(
+    exports: NewExports,
+    config?: ExportsConfig
+  ): InjectPromiseAtomApi<
+    Prettify<
+      Omit<G, 'Exports'> & {
+        Exports: (G['Exports'] extends Record<string, never>
+          ? unknown
+          : G['Exports']) &
+          NewExports
+      }
+    >,
+    EventMap,
+    Data
+  >
+
+  setExports<NewExports extends Record<string, any>>(
+    exports: NewExports,
+    config?: ExportsConfig
+  ): InjectPromiseAtomApi<
+    Prettify<Omit<G, 'Exports'> & { Exports: NewExports }>,
+    EventMap,
+    Data
+  >
+
+  setPromise<P extends Promise<any> | undefined>(
+    promise?: P
+  ): InjectPromiseAtomApi<
+    Prettify<Omit<G, 'Promise'> & { Promise: P }>,
+    EventMap,
+    Data
+  >
+
+  setTtl(
+    ttl: AtomInstanceTtl | (() => AtomInstanceTtl)
+  ): InjectPromiseAtomApi<G, EventMap, Data>
 }
 
 const hasInvalidateReason = (node: ReturnType<typeof injectSelf>) => {
@@ -111,7 +151,7 @@ export const injectPromise: {
     } & InjectSignalConfig<EventMap>
   ): InjectPromiseAtomApi<
     {
-      Exports: Record<string, any>
+      Exports: None
       Promise: Promise<Data>
       Signal: MappedSignal<{
         Events: EventMap
@@ -132,7 +172,7 @@ export const injectPromise: {
     config?: InjectPromiseConfig<Data> & InjectSignalConfig<EventMap>
   ): InjectPromiseAtomApi<
     {
-      Exports: Record<string, any>
+      Exports: None
       Promise: Promise<Data>
       Signal: MappedSignal<{
         Events: EventMap
@@ -246,7 +286,7 @@ export const injectPromise: {
     refs.current.promise
   ) as unknown as InjectPromiseAtomApi<
     {
-      Exports: Record<string, any>
+      Exports: None
       Promise: Promise<Data>
       Signal: MappedSignal<{
         Events: EventMap

packages/react/test/types.test.tsx

@@ -1031,11 +1031,33 @@ describe('react types', () => {
   })
 
   test('ResolvedState', async () => {
-    const atom1 = atom('1', () => injectPromise(() => Promise.resolve(1), []))
-    const ion1 = ion('1', () => injectPromise(() => Promise.resolve(1), []))
+    const atom1 = atom('1', () => {
+      const promiseApi = injectPromise(() => Promise.resolve(1), [])
+      const withExports = promiseApi.setExports({ a: 'a' })
+      const withTtl = withExports.setTtl(1)
+
+      // InjectPromiseAtomApi method return values retain the subtype
+      expectTypeOf(withTtl.dataSignal.get()).toEqualTypeOf<number | undefined>()
+
+      return withTtl
+    })
+
+    const ion1 = ion('1', () => {
+      const promiseApi = injectPromise(() => Promise.resolve(1), [])
+      const withExports = promiseApi.addExports({ b: 'b' })
+      const withStringPromise = withExports.setPromise(Promise.resolve('b'))
+
+      // InjectPromiseAtomApi method return values retain the subtype
+      expectTypeOf(withStringPromise.dataSignal.get()).toEqualTypeOf<
+        number | undefined
+      >()
+
+      return withStringPromise
+    })
 
     expectTypeOf<StateOf<typeof atom1>>().toEqualTypeOf<PromiseState<number>>()
     expectTypeOf(ecosystem.get(atom1).data).toEqualTypeOf<number | undefined>()
+    expectTypeOf<ExportsOf<typeof atom1>>().toEqualTypeOf<{ a: string }>()
     expectTypeOf(await ecosystem.getNode(atom1).promise).toEqualTypeOf<number>()
     expectTypeOf<ResolvedStateOf<typeof atom1>>().toEqualTypeOf<
       Omit<PromiseState<number>, 'data'> & {
@@ -1045,7 +1067,8 @@ describe('react types', () => {
 
     expectTypeOf<StateOf<typeof ion1>>().toEqualTypeOf<PromiseState<number>>()
     expectTypeOf(ecosystem.get(ion1).data).toEqualTypeOf<number | undefined>()
-    expectTypeOf(await ecosystem.getNode(ion1).promise).toEqualTypeOf<number>()
+    expectTypeOf(await ecosystem.getNode(ion1).promise).toEqualTypeOf<string>()
+    expectTypeOf<ExportsOf<typeof ion1>>().toEqualTypeOf<{ b: string }>()
     expectTypeOf<ResolvedStateOf<typeof ion1>>().toEqualTypeOf<
       Omit<PromiseState<number>, 'data'> & {
         data: number