feat: add auto hmr registration plugin to nuxt module

Author: BobbieGoede

packages/nuxt/src/auto-hmr-plugin.ts

@@ -0,0 +1,66 @@
+import type { Nuxt } from 'nuxt/schema'
+
+function getStoreDeclaration(nodes?: import('estree').VariableDeclarator[]) {
+  return nodes?.find(
+    (x) =>
+      x.init?.type === 'CallExpression' &&
+      x.init.callee.type === 'Identifier' &&
+      x.init.callee.name === 'defineStore'
+  )
+}
+
+function nameFromDeclaration(node?: import('estree').VariableDeclarator) {
+  return node?.id.type === 'Identifier' && node.id.name
+}
+
+export function autoRegisterHMRPlugin(
+  nuxt: Nuxt,
+  { resolve }: { resolve: (...path: string[]) => string }
+) {
+  const projectBasePath = resolve(nuxt.options.rootDir)
+
+  return {
+    name: 'pinia:auto-hmr-registration',
+
+    transform(code, id) {
+      if (id.startsWith('\x00')) return
+      if (!id.startsWith(projectBasePath)) return
+      if (!code.includes('defineStore') || code.includes('acceptHMRUpdate')) {
+        return
+      }
+
+      const ast = this.parse(code)
+
+      // walk top-level nodes
+      for (const n of ast.body) {
+        if (
+          n.type === 'VariableDeclaration' ||
+          n.type === 'ExportNamedDeclaration'
+        ) {
+          // find export or variable declaration that uses `defineStore`
+          const storeDeclaration = getStoreDeclaration(
+            n.type === 'VariableDeclaration'
+              ? n.declarations
+              : n.declaration?.type === 'VariableDeclaration'
+                ? n.declaration?.declarations
+                : undefined
+          )
+
+          // retrieve the variable name
+          const storeName = nameFromDeclaration(storeDeclaration)
+          if (storeName) {
+            // append HMR code
+            return {
+              code: [
+                code,
+                'if (import.meta.hot) {',
+                `  import.meta.hot.accept(acceptHMRUpdate(${storeName}, import.meta.hot))`,
+                '}',
+              ].join('\n'),
+            }
+          }
+        }
+      }
+    },
+  } satisfies import('vite').Plugin
+}

packages/nuxt/src/module.ts

@@ -7,9 +7,11 @@ import {
   addImports,
   createResolver,
   addImportsDir,
+  addVitePlugin,
 } from '@nuxt/kit'
 import type { NuxtModule } from '@nuxt/schema'
 import { fileURLToPath } from 'node:url'
+import { autoRegisterHMRPlugin } from './auto-hmr-plugin'
 
 export interface ModuleOptions {
   /**
@@ -76,6 +78,11 @@ const module: NuxtModule<ModuleOptions> = defineNuxtModule<ModuleOptions>({
         addImportsDir(resolve(nuxt.options.rootDir, storeDir))
       }
     }
+
+    // Register automatic hmr code plugin - dev mode only
+    if (nuxt.options.dev) {
+      addVitePlugin(autoRegisterHMRPlugin(nuxt, { resolve }))
+    }
   },
 })