fix(utils): unstable_unwrap to support writable atom (#1997)

* fix(utils): unwrap to support writable atom

* fix types

* fix type test

* fix type again

Author: dai-shi

src/vanilla/utils/unwrap.ts

@@ -11,6 +11,15 @@ const memo2 = <T>(create: () => T, dep1: object, dep2: object): T => {
 
 const defaultFallback = () => undefined
 
+export function unwrap<Value, Args extends unknown[], Result>(
+  anAtom: WritableAtom<Value, Args, Result>
+): WritableAtom<Awaited<Value> | undefined, Args, Result>
+
+export function unwrap<Value, Args extends unknown[], Result, PendingValue>(
+  anAtom: WritableAtom<Value, Args, Result>,
+  fallback: (prev?: Awaited<Value>) => PendingValue
+): WritableAtom<Awaited<Value> | PendingValue, Args, Result>
+
 export function unwrap<Value>(
   anAtom: Atom<Value>
 ): Atom<Awaited<Value> | undefined>
@@ -20,18 +29,18 @@ export function unwrap<Value, PendingValue>(
   fallback: (prev?: Awaited<Value>) => PendingValue
 ): Atom<Awaited<Value> | PendingValue>
 
-export function unwrap<Value, PendingValue>(
-  anAtom: Atom<Value>,
+export function unwrap<Value, Args extends unknown[], Result, PendingValue>(
+  anAtom: WritableAtom<Value, Args, Result> | Atom<Value>,
   fallback: (prev?: Awaited<Value>) => PendingValue = defaultFallback as any
 ) {
   return memo2(
     () => {
-      type PromiseAndValue = { readonly p?: Promise<Value> } & (
+      type PromiseAndValue = { readonly p?: Promise<unknown> } & (
         | { readonly v: Awaited<Value> }
         | { readonly f: PendingValue }
       )
-      const promiseErrorCache = new WeakMap<Promise<Value>, unknown>()
-      const promiseResultCache = new WeakMap<Promise<Value>, Awaited<Value>>()
+      const promiseErrorCache = new WeakMap<Promise<unknown>, unknown>()
+      const promiseResultCache = new WeakMap<Promise<unknown>, Awaited<Value>>()
       const refreshAtom = atom(0)
 
       if (import.meta.env?.MODE !== 'production') {
@@ -89,7 +98,7 @@ export function unwrap<Value, PendingValue>(
           return state.v
         }
         return state.f
-      })
+      }, (anAtom as WritableAtom<Value, unknown[], unknown>).write)
     },
     anAtom,
     fallback

tests/vanilla/utils/types.test.tsx

@@ -2,7 +2,7 @@ import { expectType } from 'ts-expect'
 import type { TypeEqual } from 'ts-expect'
 import { it } from 'vitest'
 import { atom } from 'jotai/vanilla'
-import type { Atom } from 'jotai/vanilla'
+import type { Atom, SetStateAction, WritableAtom } from 'jotai/vanilla'
 import { selectAtom, unstable_unwrap as unwrap } from 'jotai/vanilla/utils'
 
 it('selectAtom() should return the correct types', () => {
@@ -26,13 +26,28 @@ it('unwrap() should return the correct types', () => {
   const getFallbackValue = () => -1
   const syncAtom = atom(0)
   const syncUnwrappedAtom = unwrap(syncAtom, getFallbackValue)
-  expectType<TypeEqual<Atom<number>, typeof syncUnwrappedAtom>>(true)
+  expectType<
+    TypeEqual<
+      WritableAtom<number, [SetStateAction<number>], void>,
+      typeof syncUnwrappedAtom
+    >
+  >(true)
 
   const asyncAtom = atom(Promise.resolve(0))
   const asyncUnwrappedAtom = unwrap(asyncAtom, getFallbackValue)
-  expectType<TypeEqual<Atom<number>, typeof asyncUnwrappedAtom>>(true)
+  expectType<
+    TypeEqual<
+      WritableAtom<number, [SetStateAction<Promise<number>>], void>,
+      typeof asyncUnwrappedAtom
+    >
+  >(true)
 
   const maybeAsyncAtom = atom(Promise.resolve(0) as number | Promise<number>)
   const maybeAsyncUnwrappedAtom = unwrap(maybeAsyncAtom, getFallbackValue)
-  expectType<TypeEqual<Atom<number>, typeof maybeAsyncUnwrappedAtom>>(true)
+  expectType<
+    TypeEqual<
+      WritableAtom<number, [SetStateAction<number | Promise<number>>], void>,
+      typeof maybeAsyncUnwrappedAtom
+    >
+  >(true)
 })

tests/vanilla/utils/unwrap.test.ts

@@ -98,4 +98,24 @@ describe('unwrap', () => {
     store.set(countAtom, 3)
     expect(store.get(syncAtom)).toBe(3)
   })
+
+  it('should unwrap an async writable atom', async () => {
+    const store = createStore()
+    const asyncAtom = atom(Promise.resolve(1))
+    const syncAtom = unwrap(asyncAtom, (prev?: number) => prev ?? 0)
+
+    expect(store.get(syncAtom)).toBe(0)
+    await new Promise((r) => setTimeout(r)) // wait for a tick
+    expect(store.get(syncAtom)).toBe(1)
+
+    store.set(syncAtom, Promise.resolve(2))
+    expect(store.get(syncAtom)).toBe(1)
+    await new Promise((r) => setTimeout(r)) // wait for a tick
+    expect(store.get(syncAtom)).toBe(2)
+
+    store.set(syncAtom, Promise.resolve(3))
+    expect(store.get(syncAtom)).toBe(2)
+    await new Promise((r) => setTimeout(r)) // wait for a tick
+    expect(store.get(syncAtom)).toBe(3)
+  })
 })