fix(utils): Unwrap loadable/selectAtom promise types

src/utils/loadable.ts

@@ -2,12 +2,14 @@ import { atom } from 'jotai'
 import type { Atom } from 'jotai'
 import { createMemoizeAtom } from './weakCache'
 
+type ResolveType<T> = T extends Promise<infer V> ? V : T
+
 const memoizeAtom = createMemoizeAtom()
 
 type Loadable<Value> =
   | { state: 'loading' }
   | { state: 'hasError'; error: unknown }
-  | { state: 'hasData'; data: Value }
+  | { state: 'hasData'; data: ResolveType<Value> }
 
 export function loadable<Value>(anAtom: Atom<Value>): Atom<Loadable<Value>> {
   return memoizeAtom(() => {
@@ -18,7 +20,7 @@ export function loadable<Value>(anAtom: Atom<Value>): Atom<Loadable<Value>> {
       const ref = get(refAtom)
       let curr = ref.prev
       try {
-        const value = get(anAtom)
+        const value = get(anAtom) as ResolveType<Value>
         if (curr?.state !== 'hasData' || !Object.is(curr.data, value)) {
           curr = { state: 'hasData', data: value }
         }

src/utils/selectAtom.ts

@@ -2,18 +2,20 @@ import { atom } from 'jotai'
 import type { Atom } from 'jotai'
 import { createMemoizeAtom } from './weakCache'
 
+type ResolveType<T> = T extends Promise<infer V> ? V : T
+
 const memoizeAtom = createMemoizeAtom()
 
 export function selectAtom<Value, Slice>(
   anAtom: Atom<Value>,
-  selector: (v: Value) => Slice,
+  selector: (v: ResolveType<Value>) => Slice,
   equalityFn: (a: Slice, b: Slice) => boolean = Object.is
 ): Atom<Slice> {
   return memoizeAtom(() => {
     // TODO we should revisit this for a better solution than refAtom
     const refAtom = atom(() => ({} as { prev?: Slice }))
     const derivedAtom = atom((get) => {
-      const slice = selector(get(anAtom))
+      const slice = selector(get(anAtom) as ResolveType<Value>)
       const ref = get(refAtom)
       if ('prev' in ref && equalityFn(ref.prev as Slice, slice)) {
         return ref.prev as Slice

tests/utils/loadable.test.tsx

@@ -145,5 +145,8 @@ const LoadableComponent = ({ asyncAtom }: LoadableComponentProps) => {
     return <>{String(value.error)}</>
   }
 
-  return <>Data: {value.data}</>
+  // this is to ensure correct typing
+  const data: number | string = value.data
+
+  return <>Data: {data}</>
 }

tests/utils/selectAtom.test.tsx

@@ -1,4 +1,4 @@
-import { useEffect, useRef } from 'react'
+import { Suspense, useEffect, useRef } from 'react'
 import { fireEvent, render } from '@testing-library/react'
 import { atom } from 'jotai'
 import { selectAtom, useAtomValue, useUpdateAtom } from 'jotai/utils'
@@ -58,6 +58,53 @@ it('selectAtom works as expected', async () => {
   await findByText('a: 3')
 })
 
+it('selectAtom works with async atom', async () => {
+  const bigAtom = atom({ a: 0, b: 'othervalue' })
+  const bigAtomAsync = atom((get) => Promise.resolve(get(bigAtom)))
+  const littleAtom = selectAtom(bigAtomAsync, (v) => v.a)
+
+  const Parent = () => {
+    const setValue = useUpdateAtom(bigAtom)
+    return (
+      <>
+        <button
+          onClick={() =>
+            setValue((oldValue) => ({ ...oldValue, a: oldValue.a + 1 }))
+          }>
+          increment
+        </button>
+      </>
+    )
+  }
+
+  const Selector = () => {
+    const a = useAtomValue(littleAtom)
+    return (
+      <>
+        <div>a: {a}</div>
+      </>
+    )
+  }
+
+  const { findByText, getByText } = render(
+    <Provider>
+      <Suspense fallback={null}>
+        <Parent />
+        <Selector />
+      </Suspense>
+    </Provider>
+  )
+
+  await findByText('a: 0')
+
+  fireEvent.click(getByText('increment'))
+  await findByText('a: 1')
+  fireEvent.click(getByText('increment'))
+  await findByText('a: 2')
+  fireEvent.click(getByText('increment'))
+  await findByText('a: 3')
+})
+
 it('do not update unless equality function says value has changed', async () => {
   const bigAtom = atom({ a: 0 })
   const littleAtom = selectAtom(