fix(utils): Initial data for atomWithObservable

docs/utils/atom-with-observable.mdx

@@ -9,6 +9,7 @@ Ref: https://github.com/pmndrs/jotai/pull/341
 import { useAtom } from 'jotai'
 import { atomWithObservable } from 'jotai/utils'
 import { interval } from 'rxjs'
+import { map } from "rxjs/operators"
 
 const counterSubject = interval(1000).pipe(map((i) => `#${i}`));
 const counterAtom = atomWithObservable(() => counterSubject);
@@ -24,5 +25,14 @@ Its value will be last value emitted from the stream.
 
 To use this atom, you need to wrap your component with `<Suspense>`. Check out [basics/async](../basics/async.mdx).
 
+## Initial value
+`atomWithObservable` takes second optional parameter `{ initialData }` that allows to specify initial data for the atom. If `initialData` is provided then `atomWithObservable` will not suspend and will show initial value before receiving first value from observable. `initialData` can be either a value or a function that returns a value
+
+```ts
+const counterAtom = atomWithObservable(() => counterSubject, {initialData: 10});
+
+const counterAtom2 = atomWithObservable(() => counterSubject, {initialData: () => Math.random()});
+```
+
 ### Codesandbox
 <CodeSandbox id="88pnt" />

src/utils/atomWithObservable.ts

@@ -29,69 +29,52 @@ type ObservableLike<T> = {
 
 type SubjectLike<T> = ObservableLike<T> & Observer<T>
 
+type InitialDataFunction<T> = () => T | undefined
+
+export type AtomWithObservableOptions<TData> = {
+  initialData?: TData | InitialDataFunction<TData>
+}
+export type CreateObservableOptions<TData> = {
+  initialData?: TData
+}
+
 export function atomWithObservable<TData>(
-  createObservable: (get: Getter) => SubjectLike<TData>
+  createObservable: (get: Getter) => SubjectLike<TData>,
+  options?: AtomWithObservableOptions<TData>
 ): WritableAtom<TData, TData>
 
 export function atomWithObservable<TData>(
-  createObservable: (get: Getter) => ObservableLike<TData>
+  createObservable: (get: Getter) => ObservableLike<TData>,
+  options?: AtomWithObservableOptions<TData>
 ): Atom<TData>
 
 export function atomWithObservable<TData>(
-  createObservable: (get: Getter) => ObservableLike<TData> | SubjectLike<TData>
+  createObservable: (get: Getter) => ObservableLike<TData> | SubjectLike<TData>,
+  options?: AtomWithObservableOptions<TData>
 ) {
   const observableResultAtom = atom((get) => {
-    let settlePromise: ((data: TData | null, err?: unknown) => void) | null =
-      null
     let observable = createObservable(get)
     const itself = observable[Symbol.observable]?.()
     if (itself) {
       observable = itself
     }
-    const dataAtom = atom<TData | Promise<TData>>(
-      new Promise<TData>((resolve, reject) => {
-        settlePromise = (data, err) => {
-          if (err) {
-            reject(err)
-          } else {
-            resolve(data as TData)
-          }
-        }
-      })
+
+    const dataAtom = atom(
+      options?.initialData
+        ? getInitialData(options)
+        : firstValueFrom(observable)
     )
     let setData: (data: TData | Promise<TData>) => void = () => {
       throw new Error('setting data without mount')
     }
     const dataListener = (data: TData) => {
-      if (settlePromise) {
-        settlePromise(data)
-        settlePromise = null
-        if (subscription && !setData) {
-          subscription.unsubscribe()
-          subscription = null
-        }
-      } else {
-        setData(data)
-      }
+      setData(data)
     }
     const errorListener = (error: unknown) => {
-      if (settlePromise) {
-        settlePromise(null, error)
-        settlePromise = null
-        if (subscription && !setData) {
-          subscription.unsubscribe()
-          subscription = null
-        }
-      } else {
-        setData(Promise.reject<TData>(error))
-      }
+      setData(Promise.reject<TData>(error))
     }
     let subscription: Subscription | null = null
-    subscription = observable.subscribe(dataListener, errorListener)
-    if (!settlePromise) {
-      subscription.unsubscribe()
-      subscription = null
-    }
+
     dataAtom.onMount = (update) => {
       setData = update
       if (!subscription) {
@@ -107,6 +90,7 @@ export function atomWithObservable<TData>(
   const observableAtom = atom(
     (get) => {
       const { dataAtom } = get(observableResultAtom)
+
       return get(dataAtom)
     },
     (get, _set, data: TData) => {
@@ -120,3 +104,32 @@ export function atomWithObservable<TData>(
   )
   return observableAtom
 }
+
+function getInitialData<TData>(options: AtomWithObservableOptions<TData>) {
+  const initialData = options.initialData
+  return initialData instanceof Function ? initialData() : initialData
+}
+
+function firstValueFrom<T>(source: ObservableLike<T>): Promise<T> {
+  return new Promise<T>((resolve, reject) => {
+    let resolved = false
+    const subscription = source.subscribe({
+      next: (value) => {
+        resolve(value)
+        resolved = true
+        if (subscription) {
+          subscription.unsubscribe()
+        }
+      },
+      error: reject,
+      complete: () => {
+        reject()
+      },
+    })
+
+    if (resolved) {
+      // If subscription was resolved synchronously
+      subscription.unsubscribe()
+    }
+  })
+}

tests/utils/atomWithObservable.test.tsx

@@ -111,3 +111,93 @@ it('resubscribe on remount', async () => {
   act(() => subject.next(2))
   await findByText('count: 2')
 })
+
+it("count state with initialData doesn't suspend", async () => {
+  const subject = new Subject<number>()
+  const observableAtom = atomWithObservable(() => subject, { initialData: 5 })
+
+  const Counter = () => {
+    const [state] = useAtom(observableAtom)
+
+    return <>count: {state}</>
+  }
+
+  const { findByText } = render(
+    <Provider>
+      <Counter />
+    </Provider>
+  )
+
+  await findByText('count: 5')
+
+  act(() => subject.next(10))
+
+  await findByText('count: 10')
+})
+
+it('writable count state with initialData', async () => {
+  const observableAtom = atomWithObservable(
+    () => {
+      const observable = new Observable<number>((subscriber) => {
+        subscriber.next(1)
+      })
+      const subject = new Subject<number>()
+      // is this usual to delay the subscription?
+      setTimeout(() => {
+        observable.subscribe(subject)
+      }, 100)
+      return subject
+    },
+    { initialData: 5 }
+  )
+
+  const Counter = () => {
+    const [state, dispatch] = useAtom(observableAtom)
+
+    return (
+      <>
+        count: {state}
+        <button onClick={() => dispatch(9)}>button</button>
+      </>
+    )
+  }
+
+  const { findByText, getByText } = render(
+    <Provider>
+      <Suspense fallback="loading">
+        <Counter />
+      </Suspense>
+    </Provider>
+  )
+
+  await findByText('count: 5')
+
+  await findByText('count: 1')
+
+  fireEvent.click(getByText('button'))
+  await findByText('count: 9')
+})
+
+it('with initial data and synchronous subscription', async () => {
+  const observableAtom = atomWithObservable(
+    () =>
+      new Observable<number>((subscriber) => {
+        subscriber.next(1)
+      }),
+    { initialData: 5 }
+  )
+
+  const Counter = () => {
+    const [state] = useAtom(observableAtom)
+
+    return <>count: {state}</>
+  }
+
+  const { findByText } = render(
+    <Provider>
+      <Counter />
+    </Provider>
+  )
+
+  await findByText('count: 1')
+})