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

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,31 +147,42 @@ export function createJSONStorage<Value>(
       ),
     removeItem: (key) => getStringStorage()?.removeItem(key),
   }
-  if (
-    typeof window !== 'undefined' &&
-    typeof window.addEventListener === 'function' &&
-    window.Storage
-  ) {
-    storage.subscribe = (key, callback, initialValue) => {
-      if (!(getStringStorage() instanceof window.Storage)) {
-        return () => {}
-      }
-      const storageEventCallback = (e: StorageEvent) => {
-        if (e.storageArea === getStringStorage() && e.key === key) {
-          let newValue: Value
-          try {
-            newValue = JSON.parse(e.newValue || '')
-          } catch {
-            newValue = initialValue
+
+  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)
+      })
+
+  const subscriber: StringSubscribe | undefined =
+    getStringStorage()?.subscribe ||
+    (typeof window !== 'undefined' &&
+      typeof window.addEventListener === 'function' &&
+      ((key, callback) => {
+        if (!(getStringStorage() instanceof window.Storage)) {
+          return () => {}
+        }
+        const storageEventCallback = (e: StorageEvent) => {
+          if (e.storageArea === getStringStorage() && e.key === key) {
+            callback(e.newValue)
           }
-          callback(newValue)
         }
-      }
-      window.addEventListener('storage', storageEventCallback)
-      return () => {
-        window.removeEventListener('storage', storageEventCallback)
-      }
-    }
+        window.addEventListener('storage', storageEventCallback)
+        return () => {
+          window.removeEventListener('storage', storageEventCallback)
+        }
+      })) ||
+    undefined
+
+  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')
+  })
+})