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

src/vanilla/utils/atomWithStorage.ts

@@ -5,8 +5,21 @@ import { RESET } from './constants.ts'
 const isPromiseLike = (x: unknown): x is PromiseLike<unknown> =>
   typeof (x as any)?.then === 'function'
 
+type Subscribe<Value> = (
+  key: string,
+  callback: (value: Value) => void,
+  initialValue: Value,
+) => Unsubscribe
+
 type Unsubscribe = () => void
 
+type SubscribeHandler<Value> = (
+  subscribe: Subscribe<Value>,
+  key: string,
+  callback: (value: Value) => void,
+  initialValue: Value,
+) => Unsubscribe
+
 type SetStateActionWithReset<Value> =
   | Value
   | typeof RESET
@@ -38,12 +51,14 @@ export interface AsyncStringStorage {
   getItem: (key: string) => PromiseLike<string | null>
   setItem: (key: string, newValue: string) => PromiseLike<void>
   removeItem: (key: string) => PromiseLike<void>
+  subscribe?: Subscribe<string | null>
 }
 
 export interface SyncStringStorage {
   getItem: (key: string) => string | null
   setItem: (key: string, newValue: string) => void
   removeItem: (key: string) => void
+  subscribe?: Subscribe<string | null>
 }
 
 export function withStorageValidator<Value>(
@@ -114,6 +129,42 @@ export function createJSONStorage<Value>(
 ): AsyncStorage<Value> | SyncStorage<Value> {
   let lastStr: string | undefined
   let lastValue: Value
+
+  const webStorageSubscribe: Subscribe<Value> = (key, callback) => {
+    if (!(getStringStorage() instanceof window.Storage)) {
+      return () => {}
+    }
+    const storageEventCallback = (e: StorageEvent) => {
+      if (e.storageArea === getStringStorage() && e.key === key) {
+        callback((e.newValue || '') as Value)
+      }
+    }
+    window.addEventListener('storage', storageEventCallback)
+    return () => {
+      window.removeEventListener('storage', storageEventCallback)
+    }
+  }
+
+  const handleSubscribe: SubscribeHandler<Value> = (
+    subscriber,
+    key,
+    callback,
+    initialValue,
+  ) => {
+    function callbackWithParser(v: Value) {
+      let newValue: Value
+      try {
+        newValue = JSON.parse((v as string) || '')
+      } catch {
+        newValue = initialValue
+      }
+
+      callback(newValue as Value)
+    }
+
+    return subscriber(key, callbackWithParser, initialValue)
+  }
+
   const storage: AsyncStorage<Value> | SyncStorage<Value> = {
     getItem: (key, initialValue) => {
       const parse = (str: string | null) => {
@@ -141,30 +192,18 @@ export function createJSONStorage<Value>(
       ),
     removeItem: (key) => getStringStorage()?.removeItem(key),
   }
+
   if (
     typeof window !== 'undefined' &&
-    typeof window.addEventListener === 'function' &&
-    window.Storage
+    typeof window.addEventListener === 'function'
   ) {
-    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
-          }
-          callback(newValue)
-        }
-      }
-      window.addEventListener('storage', storageEventCallback)
-      return () => {
-        window.removeEventListener('storage', storageEventCallback)
-      }
+    if (getStringStorage()?.subscribe) {
+      storage.subscribe = handleSubscribe.bind(
+        null,
+        getStringStorage()!.subscribe as unknown as Subscribe<Value>,
+      )
+    } else if (window.Storage) {
+      storage.subscribe = handleSubscribe.bind(null, webStorageSubscribe)
     }
   }
   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), 1)
+  })
+
+  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')
+  })
+})