Merge pull request #321 from forestream/fix-use-theme

Fix use theme

Author: owjs3901

.changeset/short-paths-tap.md

@@ -0,0 +1,5 @@
+---
+"@devup-ui/react": patch
+---
+
+Update useTheme

packages/react/src/hooks/__tests__/use-theme.browser.test.ts

@@ -7,20 +7,26 @@ beforeEach(() => {
 describe('useTheme', () => {
   it('should return theme', async () => {
     const { useTheme } = await import('../use-theme')
-    const { result } = renderHook(() => useTheme())
+    const { result, unmount } = renderHook(() => useTheme())
     expect(result.current).toBeNull()
 
     document.documentElement.setAttribute('data-theme', 'dark')
     await waitFor(() => {
       expect(result.current).toBe('dark')
     })
-    const { result: newResult } = renderHook(() => useTheme())
+    const { result: newResult, unmount: newUnmount } = renderHook(() =>
+      useTheme(),
+    )
     expect(newResult.current).toBe('dark')
+    newUnmount()
+    unmount()
   })
+
   it('should return theme when already set', async () => {
     const { useTheme } = await import('../use-theme')
     document.documentElement.setAttribute('data-theme', 'dark')
-    const { result } = renderHook(() => useTheme())
+    const { result, unmount } = renderHook(() => useTheme())
     expect(result.current).toBe('dark')
+    unmount()
   })
 })

packages/react/src/hooks/use-theme.ts

@@ -1,48 +1,16 @@
 'use client'
-import { useId, useState } from 'react'
 
-import type { DevupTheme } from '../types/theme'
-import { useSafeEffect } from './use-safe-effect'
+import { useSyncExternalStore } from 'react'
 
-let observer: null | MutationObserver = null
-const setThemeMap: Record<string, React.Dispatch<keyof DevupTheme>> = {}
-let globalTheme: keyof DevupTheme | null = null
+import { createThemeStore } from '../stores/theme-store'
 
-export function useTheme(): keyof DevupTheme | null {
-  const id = useId()
-  const [theme, setTheme] = useState<keyof DevupTheme | null>(globalTheme)
-  useSafeEffect(() => {
-    if (globalTheme !== null) return
-    const currentTheme = document.documentElement.getAttribute('data-theme')
-    if (currentTheme !== null && currentTheme !== theme)
-      setTheme(currentTheme as keyof DevupTheme)
-  }, [theme])
-  useSafeEffect(() => {
-    const targetNode = document.documentElement
-    setThemeMap[id] = setTheme
-    if (!observer) {
-      observer = new MutationObserver(() => {
-        const theme = document.documentElement.getAttribute('data-theme')
-        globalTheme = theme as keyof DevupTheme
-        for (const key in setThemeMap)
-          setThemeMap[key](theme as keyof DevupTheme)
-      })
-      observer.observe(targetNode, {
-        attributes: true,
-        attributeFilter: ['data-theme'],
-        childList: false,
-        subtree: false,
-        characterData: false,
-        attributeOldValue: false,
-        characterDataOldValue: false,
-      })
-    }
+const themeStore = createThemeStore()
 
-    return () => {
-      delete setThemeMap[id]
-      if (observer && Object.keys(setThemeMap).length === 0)
-        observer.disconnect()
-    }
-  }, [id])
+export function useTheme() {
+  const theme = useSyncExternalStore(
+    themeStore.subscribe,
+    themeStore.get,
+    themeStore.get,
+  )
   return theme
 }

packages/react/src/stores/__tests__/theme-store.test.ts

@@ -0,0 +1,19 @@
+beforeEach(() => {
+  vi.resetModules()
+})
+
+describe('themeStore', () => {
+  it('should return themeStore object for browser', async () => {
+    const { createThemeStore } = await import('../theme-store')
+    const themeStore = createThemeStore()
+    expect(themeStore).toBeDefined()
+    expect(themeStore.get).toEqual(expect.any(Function))
+    expect(themeStore.set).toEqual(expect.any(Function))
+    expect(themeStore.subscribe).toEqual(expect.any(Function))
+    expect(themeStore.get()).toBeNull()
+    expect(themeStore.set('dark' as any)).toBeUndefined()
+    expect(themeStore.subscribe(() => {})()).toBeUndefined()
+    themeStore.subscribe(() => {})
+    expect(themeStore.set('dark' as any)).toBeUndefined()
+  })
+})

packages/react/src/stores/theme-store.ts

@@ -0,0 +1,47 @@
+'use client'
+import type { DevupTheme } from '../types/theme'
+
+type Theme = keyof DevupTheme | null
+type StoreChangeEvent = (newTheme: Theme) => void
+
+const initTheme = null
+
+export function createThemeStore() {
+  if (typeof window === 'undefined')
+    return {
+      get: () => initTheme,
+      set: () => {},
+      subscribe: () => () => {},
+    }
+
+  const el = document.documentElement
+  const subscribers: Set<StoreChangeEvent> = new Set()
+  let theme: Theme = initTheme
+  const get = () => theme
+  const set = (newTheme: Theme) => {
+    theme = newTheme
+    subscribers.forEach((subscriber) => subscriber(theme))
+  }
+
+  const subscribe = (onStoreChange: StoreChangeEvent) => {
+    subscribers.add(onStoreChange)
+    set(el.getAttribute('data-theme') as Theme)
+    return () => subscribers.delete(onStoreChange)
+  }
+
+  const mo = new MutationObserver((mutations) => {
+    for (const m of mutations)
+      if (m.type === 'attributes' && m.target instanceof HTMLElement)
+        set(m.target.getAttribute('data-theme') as Theme)
+  })
+  mo.observe(el, {
+    attributes: true,
+    attributeFilter: ['data-theme'],
+    subtree: false,
+  })
+  return {
+    get,
+    set,
+    subscribe,
+  }
+}