fix(utils): improve unstable_unwrap for sometimes async atom (pmndrs#…

…1996)

* add failing test

* fix(utils): improve unstable_unwrap for sometimes async atom

* add a test

* fix lint

* refactor

src/vanilla/utils/unwrap.ts

@@ -12,21 +12,21 @@ const memo2 = <T>(create: () => T, dep1: object, dep2: object): T => {
 const defaultFallback = () => undefined
 
 export function unwrap<Value>(
-  anAtom: Atom<Promise<Value>>
+  anAtom: Atom<Value>
 ): Atom<Awaited<Value> | undefined>
 
 export function unwrap<Value, PendingValue>(
-  anAtom: Atom<Promise<Value>>,
-  fallback: (prev?: Value) => PendingValue
+  anAtom: Atom<Value>,
+  fallback: (prev?: Awaited<Value>) => PendingValue
 ): Atom<Awaited<Value> | PendingValue>
 
 export function unwrap<Value, PendingValue>(
-  anAtom: Atom<Promise<Value>>,
-  fallback: (prev?: Value) => PendingValue = defaultFallback as any
-): Atom<Awaited<Value> | PendingValue> {
+  anAtom: Atom<Value>,
+  fallback: (prev?: Awaited<Value>) => PendingValue = defaultFallback as any
+) {
   return memo2(
     () => {
-      type PromiseAndValue = { readonly p: Promise<Value> } & (
+      type PromiseAndValue = { readonly p?: Promise<Value> } & (
         | { readonly v: Awaited<Value> }
         | { readonly f: PendingValue }
       )
@@ -45,6 +45,9 @@ export function unwrap<Value, PendingValue>(
           get(refreshAtom)
           const prev = get(promiseAndValueAtom) as PromiseAndValue | undefined
           const promise = get(anAtom)
+          if (!(promise instanceof Promise)) {
+            return { v: promise as Awaited<Value> }
+          }
           if (promise === prev?.p) {
             if (promiseErrorCache.has(promise)) {
               throw promiseErrorCache.get(promise)

tests/vanilla/utils/types.test.tsx

@@ -3,7 +3,7 @@ import type { TypeEqual } from 'ts-expect'
 import { it } from 'vitest'
 import { atom } from 'jotai/vanilla'
 import type { Atom } from 'jotai/vanilla'
-import { selectAtom } from 'jotai/vanilla/utils'
+import { selectAtom, unstable_unwrap as unwrap } from 'jotai/vanilla/utils'
 
 it('selectAtom() should return the correct types', () => {
   const doubleCount = (x: number) => x * 2
@@ -21,3 +21,18 @@ it('selectAtom() should return the correct types', () => {
     TypeEqual<Atom<number | Promise<number>>, typeof maybeAsyncSelectedAtom>
   >(true)
 })
+
+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)
+
+  const asyncAtom = atom(Promise.resolve(0))
+  const asyncUnwrappedAtom = unwrap(asyncAtom, getFallbackValue)
+  expectType<TypeEqual<Atom<number>, 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)
+})

tests/vanilla/utils/unwrap.test.ts

@@ -87,4 +87,15 @@ describe('unwrap', () => {
     await new Promise((r) => setTimeout(r)) // wait for a tick
     expect(store.get(syncAtom)).toBe(6)
   })
+
+  it('should unwrap a sync atom which is noop', async () => {
+    const store = createStore()
+    const countAtom = atom(1)
+    const syncAtom = unwrap(countAtom)
+    expect(store.get(syncAtom)).toBe(1)
+    store.set(countAtom, 2)
+    expect(store.get(syncAtom)).toBe(2)
+    store.set(countAtom, 3)
+    expect(store.get(syncAtom)).toBe(3)
+  })
 })