feat(utils): allow providing a subscribe method to createJSONStorage …

…util (#2539)

* feat: allow providing a subscribe method to createJSONStorage util

* refactor: convert handleSubscribe to a HOF

* refactor

* refactor 2

---------

Co-authored-by: Daishi Kato <dai-shi@users.noreply.github.com>
Co-authored-by: daishi <daishi@axlight.com>

src/vanilla/utils/atomWithStorage.ts

@@ -7,6 +7,17 @@ const isPromiseLike = (x: unknown): x is PromiseLike<unknown> =>
 
 type Unsubscribe = () => void
 
+type Subscribe<Value> = (
+  key: string,
+  callback: (value: Value) => void,
+  initialValue: Value,
+) => Unsubscribe
+
+type StringSubscribe = (
+  key: string,
+  callback: (value: string | null) => void,
+) => Unsubscribe
+
 type SetStateActionWithReset<Value> =
   | Value
   | typeof RESET
@@ -16,34 +27,28 @@ export interface AsyncStorage<Value> {
   getItem: (key: string, initialValue: Value) => PromiseLike<Value>
   setItem: (key: string, newValue: Value) => PromiseLike<void>
   removeItem: (key: string) => PromiseLike<void>
-  subscribe?: (
-    key: string,
-    callback: (value: Value) => void,
-    initialValue: Value,
-  ) => Unsubscribe
+  subscribe?: Subscribe<Value>
 }
 
 export interface SyncStorage<Value> {
   getItem: (key: string, initialValue: Value) => Value
   setItem: (key: string, newValue: Value) => void
   removeItem: (key: string) => void
-  subscribe?: (
-    key: string,
-    callback: (value: Value) => void,
-    initialValue: Value,
-  ) => Unsubscribe
+  subscribe?: Subscribe<Value>
 }
 
 export interface AsyncStringStorage {
   getItem: (key: string) => PromiseLike<string | null>
   setItem: (key: string, newValue: string) => PromiseLike<void>
   removeItem: (key: string) => PromiseLike<void>
+  subscribe?: StringSubscribe
 }
 
 export interface SyncStringStorage {
   getItem: (key: string) => string | null
   setItem: (key: string, newValue: string) => void
   removeItem: (key: string) => void
+  subscribe?: StringSubscribe
 }
 
 export function withStorageValidator<Value>(
@@ -114,6 +119,7 @@ export function createJSONStorage<Value>(
 ): AsyncStorage<Value> | SyncStorage<Value> {
   let lastStr: string | undefined
   let lastValue: Value
+
   const storage: AsyncStorage<Value> | SyncStorage<Value> = {
     getItem: (key, initialValue) => {
       const parse = (str: string | null) => {
@@ -141,24 +147,32 @@ export function createJSONStorage<Value>(
       ),
     removeItem: (key) => getStringStorage()?.removeItem(key),
   }
+
+  const createHandleSubscribe =
+    (subscriber: StringSubscribe): Subscribe<Value> =>
+    (key, callback, initialValue) =>
+      subscriber(key, (v) => {
+        let newValue: Value
+        try {
+          newValue = JSON.parse(v || '')
+        } catch {
+          newValue = initialValue
+        }
+        callback(newValue)
+      })
+
+  let subscriber = getStringStorage()?.subscribe
   if (
+    !subscriber &&
     typeof window !== 'undefined' &&
     typeof window.addEventListener === 'function' &&
-    window.Storage
+    window.Storage &&
+    getStringStorage() instanceof window.Storage
   ) {
-    storage.subscribe = (key, callback, initialValue) => {
-      if (!(getStringStorage() instanceof window.Storage)) {
-        return () => {}
-      }
+    subscriber = (key, callback) => {
       const storageEventCallback = (e: StorageEvent) => {
         if (e.storageArea === getStringStorage() && e.key === key) {
-          let newValue: Value
-          try {
-            newValue = JSON.parse(e.newValue || '')
-          } catch {
-            newValue = initialValue
-          }
-          callback(newValue)
+          callback(e.newValue)
         }
       }
       window.addEventListener('storage', storageEventCallback)
@@ -167,6 +181,10 @@ export function createJSONStorage<Value>(
       }
     }
   }
+
+  if (subscriber) {
+    storage.subscribe = createHandleSubscribe(subscriber)
+  }
   return storage
 }
 

tests/react/vanilla-utils/atomWithStorage.test.tsx

@@ -9,6 +9,7 @@ import {
   createJSONStorage,
   unstable_withStorageValidator as withStorageValidator,
 } from 'jotai/vanilla/utils'
+import type { SyncStringStorage } from 'jotai/vanilla/utils/atomWithStorage'
 
 const resolve: (() => void)[] = []
 
@@ -601,3 +602,101 @@ describe('withStorageValidator', () => {
     atomWithStorage('my-number', 0, withStorageValidator(isNumber)(storage))
   })
 })
+
+describe('with subscribe method in string storage', () => {
+  it('createJSONStorage subscriber is called correctly', async () => {
+    const store = createStore()
+
+    const subscribe = vi.fn()
+    const stringStorage = {
+      getItem: () => {
+        return null
+      },
+      setItem: () => {},
+      removeItem: () => {},
+      subscribe,
+    }
+
+    const dummyStorage = createJSONStorage<number>(() => stringStorage)
+
+    const dummyAtom = atomWithStorage<number>('dummy', 1, dummyStorage)
+
+    const DummyComponent = () => {
+      const [value] = useAtom(dummyAtom, { store })
+      return (
+        <>
+          <div>{value}</div>
+        </>
+      )
+    }
+
+    render(
+      <StrictMode>
+        <DummyComponent />
+      </StrictMode>,
+    )
+
+    expect(subscribe).toHaveBeenCalledWith('dummy', expect.any(Function))
+  })
+
+  it('createJSONStorage subscriber responds to events correctly', async () => {
+    const storageData: Record<string, string> = {
+      count: '10',
+    }
+
+    const stringStorage = {
+      getItem: (key: string) => {
+        return storageData[key] || null
+      },
+      setItem: (key: string, newValue: string) => {
+        storageData[key] = newValue
+      },
+      removeItem: (key: string) => {
+        delete storageData[key]
+      },
+      subscribe(key, callback) {
+        function handler(event: CustomEvent<string>) {
+          callback(event.detail)
+        }
+
+        window.addEventListener('dummystoragechange', handler as EventListener)
+        return () =>
+          window.removeEventListener(
+            'dummystoragechange',
+            handler as EventListener,
+          )
+      },
+    } as SyncStringStorage
+
+    const dummyStorage = createJSONStorage<number>(() => stringStorage)
+
+    const countAtom = atomWithStorage('count', 1, dummyStorage)
+
+    const Counter = () => {
+      const [count] = useAtom(countAtom)
+      return (
+        <>
+          <div>count: {count}</div>
+        </>
+      )
+    }
+
+    const { findByText } = render(
+      <StrictMode>
+        <Counter />
+      </StrictMode>,
+    )
+
+    await findByText('count: 10')
+
+    storageData.count = '12'
+    fireEvent(
+      window,
+      new CustomEvent('dummystoragechange', {
+        detail: '12',
+      }),
+    )
+    await findByText('count: 12')
+    // expect(storageData.count).toBe('11')
+  })
+})