From 2c99bb80c72fdbde9cd2ff3ad7caae0be632b864 Mon Sep 17 00:00:00 2001
From: Benjamin Canac <canacb1@gmail.com>
Date: Wed, 18 Dec 2024 12:45:51 +0100
Subject: [PATCH] fix(templates): infer variants types in generated theme

---
 src/templates.ts | 42 ++++++++++++++++++++++++++++++------------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/src/templates.ts b/src/templates.ts
index cf247f98..6a41418e 100644
--- a/src/templates.ts
+++ b/src/templates.ts
@@ -24,29 +24,47 @@ export function getTemplates(options: ModuleOptions, uiConfig: Record<string, an
         const template = (theme as any)[component]
         const result = typeof template === 'function' ? template(options) : template
 
-        const variants = Object.keys(result.variants || {})
+        const variants = Object.entries(result.variants || {})
+          .filter(([_, values]) => {
+            const keys = Object.keys(values as Record<string, unknown>)
+            return keys.some(key => key !== 'true' && key !== 'false')
+          })
+          .map(([key]) => key)
 
         let json = JSON.stringify(result, null, 2)
 
         for (const variant of variants) {
-          json = json.replace(new RegExp(`("${variant}": "[^"]+")`, 'g'), '$1 as const')
+          json = json.replace(new RegExp(`("${variant}": "[^"]+")`, 'g'), `$1 as typeof ${variant}[number]`)
           json = json.replace(new RegExp(`("${variant}": \\[\\s*)((?:"[^"]+",?\\s*)+)(\\])`, 'g'), (_, before, match, after) => {
-            const replaced = match.replace(/("[^"]+")/g, '$1 as const')
+            const replaced = match.replace(/("[^"]+")/g, `$1 as typeof ${variant}[number]`)
             return `${before}${replaced}${after}`
           })
         }
 
-        // For local development, directly import from theme
-        if (process.env.DEV) {
-          return [
-            `import template from ${JSON.stringify(fileURLToPath(new URL(`./theme/${kebabCase(component)}`, import.meta.url)))}`,
-            `const result = typeof template === 'function' ? template(${JSON.stringify(options)}) : template`,
-            `const json = ${json}`,
-            `export default result as typeof json`
-          ].join('\n')
+        function generateVariantDeclarations(variants: string[]) {
+          return variants.map((variant) => {
+            const keys = Object.keys(result.variants[variant])
+            return `const ${variant} = ${JSON.stringify(keys, null, 2)} as const`
+          })
         }
 
-        return `export default ${json}`
+        // For local development, import directly from theme
+        if (process.env.DEV) {
+          const templatePath = fileURLToPath(new URL(`./theme/${kebabCase(component)}`, import.meta.url))
+          return [
+            `import template from ${JSON.stringify(templatePath)}`,
+            ...generateVariantDeclarations(variants),
+            `const result = typeof template === 'function' ? template(${JSON.stringify(options, null, 2)}) : template`,
+            `const theme = ${json}`,
+            `export default result as typeof theme`
+          ].join('\n\n')
+        }
+
+        // For production build
+        return [
+          ...generateVariantDeclarations(variants),
+          `export default ${json}`
+        ].join('\n\n')
       }
     })
   }