Port #178 (#200)

* port pr 178 * fix up logic
pacocoursey · Aug 18, 2023 · f0668bd · f0668bd
1 parent aba0449
commit f0668bd
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 15 deletions.
diff --git a/packages/next-themes/src/index.tsx b/packages/next-themes/src/index.tsx
@@ -26,7 +26,7 @@ export const ThemeProvider: React.FC<ThemeProviderProps> = props => {
   return <Theme {...props} />
 }
 
-const defaultThemes = ['light', 'dark'];
+const defaultThemes = ['light', 'dark']
 
 const Theme: React.FC<ThemeProviderProps> = ({
   forcedTheme,
@@ -82,11 +82,12 @@ const Theme: React.FC<ThemeProviderProps> = ({
 
   const setTheme = useCallback(
     theme => {
-      setThemeState(theme)
+      const newTheme = typeof theme === 'function' ? theme(theme) : theme
+      setThemeState(newTheme)
 
       // Save to storage
       try {
-        localStorage.setItem(storageKey, theme)
+        localStorage.setItem(storageKey, newTheme)
       } catch (e) {
         // Unsupported
       }
@@ -138,19 +139,20 @@ const Theme: React.FC<ThemeProviderProps> = ({
     applyTheme(forcedTheme ?? theme)
   }, [forcedTheme, theme])
 
-  const providerValue = useMemo(() => ({
-    theme,
-    setTheme,
-    forcedTheme,
-    resolvedTheme: theme === 'system' ? resolvedTheme : theme,
-    themes: enableSystem ? [...themes, 'system'] : themes,
-    systemTheme: (enableSystem ? resolvedTheme : undefined) as 'light' | 'dark' | undefined
-  }), [theme, setTheme, forcedTheme, resolvedTheme, enableSystem, themes]);
+  const providerValue = useMemo(
+    () => ({
+      theme,
+      setTheme,
+      forcedTheme,
+      resolvedTheme: theme === 'system' ? resolvedTheme : theme,
+      themes: enableSystem ? [...themes, 'system'] : themes,
+      systemTheme: (enableSystem ? resolvedTheme : undefined) as 'light' | 'dark' | undefined
+    }),
+    [theme, setTheme, forcedTheme, resolvedTheme, enableSystem, themes]
+  )
 
   return (
-    <ThemeContext.Provider
-      value={providerValue}
-    >
+    <ThemeContext.Provider value={providerValue}>
       <ThemeScript
         {...{
           forcedTheme,

diff --git a/packages/next-themes/src/types.ts b/packages/next-themes/src/types.ts
@@ -8,7 +8,7 @@ export interface UseThemeProps {
   /** Forced theme name for the current page */
   forcedTheme?: string
   /** Update the theme */
-  setTheme: (theme: string) => void
+  setTheme: React.Dispatch<React.SetStateAction<string>>
   /** Active theme name */
   theme?: string
   /** If `enableSystem` is true and the active theme is "system", this returns whether the system preference resolved to "dark" or "light". Otherwise, identical to `theme` */