fix: allow multiple attributes in the injection script

pacocoursey · Aug 31, 2024 · a4589b0 · a4589b0
1 parent bf0c5a4
commit a4589b0
Showing 1 changed file with 12 additions and 8 deletions.
diff --git a/next-themes/src/script.ts b/next-themes/src/script.ts
@@ -10,16 +10,20 @@ export const script = (
 ) => {
   const el = document.documentElement
   const systemThemes = ['light', 'dark']
-  const isClass = attribute === 'class'
-  const classes = isClass && value ? themes.map(t => value[t] || t) : themes
 
   function updateDOM(theme: string) {
-    if (isClass) {
-      el.classList.remove(...classes)
-      el.classList.add(theme)
-    } else {
-      el.setAttribute(attribute, theme)
-    }
+    const attributes = Array.isArray(attribute) ? attribute : [attribute]
+
+    attributes.forEach(attr => {
+      const isClass = attr === 'class'
+      const classes = isClass && value ? themes.map(t => value[t] || t) : themes
+      if (isClass) {
+        el.classList.remove(...classes)
+        el.classList.add(theme)
+      } else {
+        el.setAttribute(attr, theme)
+      }
+    })
 
     setColorScheme(theme)
   }