Add write to atomWithDefault

pmndrs · Mar 26, 2021 · 3320848 · 3320848
1 parent 39f6fc0
commit 3320848
Show file tree

Hide file tree

Showing 3 changed files with 119 additions and 17 deletions.
diff --git a/.size-snapshot.json b/.size-snapshot.json
@@ -1,8 +1,8 @@
 {
   "index.module.js": {
-    "bundled": 19949,
-    "minified": 8342,
-    "gzipped": 2804,
+    "bundled": 19955,
+    "minified": 8345,
+    "gzipped": 2805,
     "treeshaked": {
       "rollup": {
         "code": 14,
@@ -14,9 +14,9 @@
     }
   },
   "utils.module.js": {
-    "bundled": 7958,
-    "minified": 3528,
-    "gzipped": 1427,
+    "bundled": 8383,
+    "minified": 3692,
+    "gzipped": 1489,
     "treeshaked": {
       "rollup": {
         "code": 28,
@@ -98,14 +98,14 @@
     }
   },
   "index.js": {
-    "bundled": 23392,
-    "minified": 10132,
-    "gzipped": 3391
+    "bundled": 23398,
+    "minified": 10135,
+    "gzipped": 3392
   },
   "utils.js": {
-    "bundled": 11093,
-    "minified": 5420,
-    "gzipped": 2040
+    "bundled": 11566,
+    "minified": 5645,
+    "gzipped": 2105
   },
   "devtools.js": {
     "bundled": 2394,

diff --git a/src/utils/atomWithDefault.ts b/src/utils/atomWithDefault.ts
@@ -1,11 +1,21 @@
 import { atom, PrimitiveAtom } from 'jotai'
-import type { Read } from '../core/types'
+import type { Read, Write, WritableAtom } from '../core/types'
 
 export function atomWithDefault<Value>(
   getDefault: Read<Value>
-): PrimitiveAtom<Value> {
+): PrimitiveAtom<Value>
+
+export function atomWithDefault<Value, Update>(
+  getDefault: Read<Value>,
+  write: Write<Update>
+): WritableAtom<Value, Update>
+
+export function atomWithDefault<Value, Update>(
+  getDefault: Read<Value>,
+  write?: Write<Update>
+): WritableAtom<Value, Update> {
   const overwrittenAtom = atom(false)
-  const anAtom: PrimitiveAtom<Value> = atom(
+  const anAtom: any = atom<Value, Update>(
     (get) => {
       if (get(overwrittenAtom)) {
         return get(anAtom)
@@ -16,8 +26,10 @@ export function atomWithDefault<Value>(
       set(overwrittenAtom, true)
       set(
         anAtom,
-        typeof update === 'function'
-          ? (update as (prev: Value) => Value)(get(anAtom))
+        write
+          ? write(get, set, update)
+          : typeof update === 'function'
+          ? update(get(anAtom))
           : update
       )
     }

diff --git a/tests/utils/atomWithDefault.test.tsx b/tests/utils/atomWithDefault.test.tsx
@@ -41,6 +41,47 @@ it('simple sync get default', async () => {
   await findByText('count1: 3, count2: 5')
 })
 
+it('simple sync set default', async () => {
+  const count1Atom = atom(1)
+  const count2Atom = atomWithDefault(
+    (get) => get(count1Atom) * 2,
+    (get, set) => {
+      set(count2Atom, get(count1Atom) * 5)
+    }
+  )
+
+  const Counter: React.FC = () => {
+    const [count1, setCount1] = useAtom(count1Atom)
+    const [count2, setCount2] = useAtom(count2Atom)
+    return (
+      <>
+        <div>
+          count1: {count1}, count2: {count2}
+        </div>
+        <button onClick={() => setCount1((c) => c + 1)}>button1</button>
+        <button onClick={() => setCount2('asd')}>button2</button>
+      </>
+    )
+  }
+
+  const { findByText, getByText } = render(
+    <Provider>
+      <Counter />
+    </Provider>
+  )
+
+  await findByText('count1: 1, count2: 2')
+
+  fireEvent.click(getByText('button1'))
+  await findByText('count1: 2, count2: 4')
+
+  fireEvent.click(getByText('button2'))
+  await findByText('count1: 2, count2: 10')
+
+  fireEvent.click(getByText('button1'))
+  await findByText('count1: 3, count2: 10')
+})
+
 it('simple async get default', async () => {
   const count1Atom = atom(1)
   const count2Atom = atomWithDefault(async (get) => {
@@ -83,3 +124,52 @@ it('simple async get default', async () => {
   fireEvent.click(getByText('button1'))
   await findByText('count1: 3, count2: 5')
 })
+
+it('simple async set default', async () => {
+  const count1Atom = atom(1)
+  const count2Atom = atomWithDefault(
+    async (get) => {
+      await new Promise((r) => setTimeout(r, 10))
+      return get(count1Atom) * 2
+    },
+    async (get, set, update: number) => {
+      await new Promise((r) => setTimeout(r, 10))
+      set(count2Atom, get(count1Atom) * update)
+    }
+  )
+
+  const Counter: React.FC = () => {
+    const [count1, setCount1] = useAtom(count1Atom)
+    const [count2, setCount2] = useAtom(count2Atom)
+    return (
+      <>
+        <div>
+          count1: {count1}, count2: {count2}
+        </div>
+        <button onClick={() => setCount1((c) => c + 1)}>button1</button>
+        <button onClick={() => setCount2(5)}>button2</button>
+      </>
+    )
+  }
+
+  const { findByText, getByText } = render(
+    <Provider>
+      <Suspense fallback="loading">
+        <Counter />
+      </Suspense>
+    </Provider>
+  )
+
+  await findByText('loading')
+  await findByText('count1: 1, count2: 2')
+
+  fireEvent.click(getByText('button1'))
+  await findByText('loading')
+  await findByText('count1: 2, count2: 4')
+
+  fireEvent.click(getByText('button2'))
+  await findByText('count1: 2, count2: 10')
+
+  fireEvent.click(getByText('button1'))
+  await findByText('count1: 3, count2: 10')
+})