fix: support HMR for class components (#320)

Author: ArnaudBarre

packages/plugin-react/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Support HMR for class components
+
+This is a long overdue and should fix some issues people had with HMR when migrating from CRA.
+
 ## 4.2.1 (2023-12-04)
 
 Remove generic parameter on `Plugin` to avoid type error with Rollup 4/Vite 5 and `skipLibCheck: false`.

packages/plugin-react/src/fast-refresh.ts

@@ -28,10 +28,12 @@ window.$RefreshSig$ = () => (type) => type
 window.__vite_plugin_react_preamble_installed__ = true
 `
 
-const header = `
+const sharedHeader = `
 import RefreshRuntime from "${runtimePublicPath}";
 
 const inWebWorker = typeof WorkerGlobalScope !== 'undefined' && self instanceof WorkerGlobalScope;
+`.replace(/\n+/g, '')
+const functionHeader = `
 let prevRefreshReg;
 let prevRefreshSig;
 
@@ -51,11 +53,13 @@ if (import.meta.hot && !inWebWorker) {
   window.$RefreshSig$ = RefreshRuntime.createSignatureFunctionForTransform;
 }`.replace(/\n+/g, '')
 
-const footer = `
+const functionFooter = `
 if (import.meta.hot && !inWebWorker) {
   window.$RefreshReg$ = prevRefreshReg;
   window.$RefreshSig$ = prevRefreshSig;
-
+}`
+const sharedFooter = `
+if (import.meta.hot && !inWebWorker) {
   RefreshRuntime.__hmr_import(import.meta.url).then((currentExports) => {
     RefreshRuntime.registerExportsForReactRefresh(__SOURCE__, currentExports);
     import.meta.hot.accept((nextExports) => {
@@ -68,8 +72,19 @@ if (import.meta.hot && !inWebWorker) {
 
 export function addRefreshWrapper(code: string, id: string): string {
   return (
-    header.replace('__SOURCE__', JSON.stringify(id)) +
+    sharedHeader +
+    functionHeader.replace('__SOURCE__', JSON.stringify(id)) +
     code +
-    footer.replace('__SOURCE__', JSON.stringify(id))
+    functionFooter +
+    sharedFooter.replace('__SOURCE__', JSON.stringify(id))
+  )
+}
+
+export function addClassComponentRefreshWrapper(
+  code: string,
+  id: string,
+): string {
+  return (
+    sharedHeader + code + sharedFooter.replace('__SOURCE__', JSON.stringify(id))
   )
 }

packages/plugin-react/src/index.ts

@@ -11,6 +11,7 @@ import type {
   UserConfig,
 } from 'vite'
 import {
+  addClassComponentRefreshWrapper,
   addRefreshWrapper,
   preambleCode,
   runtimeCode,
@@ -86,6 +87,7 @@ export type ViteReactPluginApi = {
   reactBabel?: ReactBabelHook
 }
 
+const reactCompRE = /extends\s+(?:React\.)?(?:Pure)?Component/
 const refreshContentRE = /\$Refresh(?:Reg|Sig)\$\(/
 const defaultIncludeRE = /\.[tj]sx?$/
 const tsRE = /\.tsx?$/
@@ -250,8 +252,12 @@ export default function viteReact(opts: Options = {}): PluginOption[] {
 
       if (result) {
         let code = result.code!
-        if (useFastRefresh && refreshContentRE.test(code)) {
-          code = addRefreshWrapper(code, id)
+        if (useFastRefresh) {
+          if (refreshContentRE.test(code)) {
+            code = addRefreshWrapper(code, id)
+          } else if (reactCompRE.test(code)) {
+            code = addClassComponentRefreshWrapper(code, id)
+          }
         }
         return { code, map: result.map }
       }

playground/class-components/__tests__/class-components.spec.ts

@@ -0,0 +1,30 @@
+import { expect, test } from 'vitest'
+import {
+  editFile,
+  isServe,
+  page,
+  untilBrowserLogAfter,
+  untilUpdated,
+} from '~utils'
+
+test('should render', async () => {
+  expect(await page.textContent('span')).toMatch('Hello World')
+})
+
+if (isServe) {
+  test('Class component HMR', async () => {
+    editFile('src/App.tsx', (code) => code.replace('World', 'class components'))
+    await untilBrowserLogAfter(
+      () => page.textContent('span'),
+      '[vite] hot updated: /src/App.tsx',
+    )
+    await untilUpdated(() => page.textContent('span'), 'Hello class components')
+
+    editFile('src/utils.tsx', (code) => code.replace('Hello', 'Hi'))
+    await untilBrowserLogAfter(
+      () => page.textContent('span'),
+      '[vite] hot updated: /src/App.tsx',
+    )
+    await untilUpdated(() => page.textContent('span'), 'Hi class components')
+  })
+}

playground/class-components/index.html

@@ -0,0 +1,13 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <link rel="icon" type="image/svg+xml" href="/vite.svg" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>Vite + React + class components</title>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/index.tsx"></script>
+  </body>
+</html>

playground/class-components/package.json

@@ -0,0 +1,19 @@
+{
+  "name": "@vitejs/test-class-components",
+  "private": true,
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "vite build",
+    "preview": "vite preview"
+  },
+  "dependencies": {
+    "react": "^18.3.1",
+    "react-dom": "^18.3.1"
+  },
+  "devDependencies": {
+    "@types/react": "^18.3.2",
+    "@types/react-dom": "^18.3.0",
+    "@vitejs/plugin-react": "workspace:*"
+  }
+}

playground/class-components/public/vite.svg

@@ -0,0 +1,8 @@
+import { Component } from 'react'
+import { getGetting } from './utils'
+
+export class App extends Component {
+  render() {
+    return <span>{getGetting()} World</span>
+  }
+}

playground/class-components/src/App.tsx

@@ -0,0 +1,9 @@
+import { StrictMode } from 'react'
+import { createRoot } from 'react-dom/client'
+import { App } from './App'
+
+createRoot(document.getElementById('root')!).render(
+  <StrictMode>
+    <App />
+  </StrictMode>,
+)

playground/class-components/src/index.tsx

@@ -0,0 +1 @@
+export const getGetting = () => <span>Hello</span>