Vite: Improve handling of preview annotations

code/builders/builder-vite/package.json

@@ -50,14 +50,16 @@
     "express": "^4.19.2",
     "find-cache-dir": "^3.0.0",
     "fs-extra": "^11.1.0",
+    "knitwork": "^1.1.0",
     "magic-string": "^0.30.0",
+    "pathe": "^1.1.2",
+    "slash": "^5.0.0",
     "ts-dedent": "^2.0.0"
   },
   "devDependencies": {
     "@types/express": "^4.17.21",
     "@types/node": "^18.0.0",
     "glob": "^10.0.0",
-    "slash": "^5.0.0",
     "typescript": "^5.3.2",
     "vite": "^4.0.4"
   },

code/builders/builder-vite/src/__snapshots__/codegen-modern-iframe-script.test.ts.snap

@@ -0,0 +1,61 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`generateModernIframeScriptCodeFromPreviews > handle multiple annotations 1`] = `
+"import { composeConfigs, PreviewWeb, ClientApi } from 'storybook/internal/preview-api';
+import '/virtual:/@storybook/builder-vite/setup-addons.js';
+import { importFn } from '/virtual:/@storybook/builder-vite/storybook-stories.js';
+import previewAnnotations1_2032 from "/previewAnnotations1";
+import previewAnnotations2_2033 from "/previewAnnotations2";
+const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+  const configs = [
+    hmrPreviewAnnotationModules.at(0) ?? previewAnnotations1_2032,
+    hmrPreviewAnnotationModules.at(1) ?? previewAnnotations2_2033
+  ]
+  return composeConfigs(configs);
+}
+
+window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
+
+window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
+
+if (import.meta.hot) {
+  import.meta.hot.accept('/virtual:/@storybook/builder-vite/storybook-stories.js', (newModule) => {
+    // importFn has changed so we need to patch the new one in
+    window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
+  });
+
+  import.meta.hot.accept(["/previewAnnotations1","/previewAnnotations2"], (previewAnnotationModules) => {
+    // getProjectAnnotations has changed so we need to patch the new one in
+    window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
+  });
+};"
+`;
+
+exports[`generateModernIframeScriptCodeFromPreviews > handle one annotation 1`] = `
+"import { composeConfigs, PreviewWeb, ClientApi } from 'storybook/internal/preview-api';
+import '/virtual:/@storybook/builder-vite/setup-addons.js';
+import { importFn } from '/virtual:/@storybook/builder-vite/storybook-stories.js';
+import previewAnnotations_1983 from "/previewAnnotations";
+const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+  const configs = [
+    hmrPreviewAnnotationModules.at(0) ?? previewAnnotations_1983
+  ]
+  return composeConfigs(configs);
+}
+
+window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
+
+window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
+
+if (import.meta.hot) {
+  import.meta.hot.accept('/virtual:/@storybook/builder-vite/storybook-stories.js', (newModule) => {
+    // importFn has changed so we need to patch the new one in
+    window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
+  });
+
+  import.meta.hot.accept(["/previewAnnotations"], (previewAnnotationModules) => {
+    // getProjectAnnotations has changed so we need to patch the new one in
+    window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
+  });
+};"
+`;

code/builders/builder-vite/src/codegen-modern-iframe-script.test.ts

@@ -0,0 +1,23 @@
+import { describe, it, expect } from 'vitest';
+import { generateModernIframeScriptCodeFromPreviews } from './codegen-modern-iframe-script';
+
+const projectRoot = 'projectRoot';
+
+describe('generateModernIframeScriptCodeFromPreviews', () => {
+  it('handle one annotation', async () => {
+    const result = await generateModernIframeScriptCodeFromPreviews({
+      previewAnnotations: ['previewAnnotations'],
+      projectRoot,
+      frameworkName: 'frameworkName',
+    });
+    expect(result).toMatchSnapshot();
+  });
+  it('handle multiple annotations', async () => {
+    const result = await generateModernIframeScriptCodeFromPreviews({
+      previewAnnotations: ['previewAnnotations1', 'previewAnnotations2'],
+      projectRoot,
+      frameworkName: 'frameworkName',
+    });
+    expect(result).toMatchSnapshot();
+  });
+});

code/builders/builder-vite/src/codegen-modern-iframe-script.ts

@@ -3,6 +3,10 @@ import type { Options, PreviewAnnotation } from 'storybook/internal/types';
 import { virtualStoriesFile, virtualAddonSetupFile } from './virtual-file-names';
 import { processPreviewAnnotation } from './utils/process-preview-annotation';
 
+import { genArrayFromRaw, genImport, genSafeVariableName } from 'knitwork';
+import { filename } from 'pathe/utils';
+import { dedent } from 'ts-dedent';
+
 export async function generateModernIframeScriptCode(options: Options, projectRoot: string) {
   const { presets, configDir } = options;
   const frameworkName = await getFrameworkName(options);
@@ -13,27 +17,52 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
     [],
     options
   );
-  const previewAnnotationURLs = [...previewAnnotations, previewOrConfigFile]
-    .filter(Boolean)
+  return generateModernIframeScriptCodeFromPreviews({
+    previewAnnotations: [...previewAnnotations, previewOrConfigFile],
+    projectRoot,
+    frameworkName,
+  });
+}
+
+export async function generateModernIframeScriptCodeFromPreviews(options: {
+  previewAnnotations: (string | undefined)[];
+  projectRoot: string;
+  frameworkName: string;
+}) {
+  const { projectRoot, frameworkName } = options;
+  const previewAnnotationURLs = options.previewAnnotations
+    .filter((path) => path !== undefined)
     .map((path) => processPreviewAnnotation(path, projectRoot));
 
+  const variables: string[] = [];
+  const imports: string[] = [];
+  for (const previewAnnotation of previewAnnotationURLs) {
+    const variable =
+      genSafeVariableName(filename(previewAnnotation)).replace(/_(45|46|47)/g, '_') +
+      '_' +
+      hash(previewAnnotation);
+    variables.push(variable);
+    imports.push(genImport(previewAnnotation, variable));
+  }
+
   // This is pulled out to a variable because it is reused in both the initial page load
-  // and the HMR handler.  We don't use the hot.accept callback params because only the changed
-  // modules are provided, the rest are null.  We can just re-import everything again in that case.
-  const getPreviewAnnotationsFunction = `
-  const getProjectAnnotations = async (hmrPreviewAnnotationModules = []) => {
-    const configs = await Promise.all([${previewAnnotationURLs
-      .map(
+  // and the HMR handler.
+  // The `hmrPreviewAnnotationModules` parameter is used to pass the updated modules from HMR.
+  // However, only the changed modules are provided, the rest are null.
+  const getPreviewAnnotationsFunction = dedent`
+  const getProjectAnnotations = (hmrPreviewAnnotationModules = []) => {
+    const configs = ${genArrayFromRaw(
+      variables.map(
         (previewAnnotation, index) =>
-          // Prefer the updated module from an HMR update, otherwise import the original module
-          `hmrPreviewAnnotationModules.at(${index}) ?? import('${previewAnnotation}')`
-      )
-      .join(',\n')}])
+          // Prefer the updated module from an HMR update, otherwise the original module
+          `hmrPreviewAnnotationModules.at(${index}) ?? ${previewAnnotation}`
+      ),
+      '  '
+    )}
     return composeConfigs(configs);
   }`;
 
-  // eslint-disable-next-line @typescript-eslint/no-shadow
-  const generateHMRHandler = (frameworkName: string): string => {
+  const generateHMRHandler = (): string => {
     // Web components are not compatible with HMR, so disable HMR, reload page instead.
     if (frameworkName === '@storybook/web-components-vite') {
       return `
@@ -42,19 +71,18 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
       }`.trim();
     }
 
-    return `
+    return dedent`
     if (import.meta.hot) {
       import.meta.hot.accept('${virtualStoriesFile}', (newModule) => {
-      // importFn has changed so we need to patch the new one in
-      window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
+        // importFn has changed so we need to patch the new one in
+        window.__STORYBOOK_PREVIEW__.onStoriesChanged({ importFn: newModule.importFn });
       });
 
-    import.meta.hot.accept(${JSON.stringify(previewAnnotationURLs)}, (previewAnnotationModules) => {
-      ${getPreviewAnnotationsFunction}
-      // getProjectAnnotations has changed so we need to patch the new one in
-      window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
-    });
-  }`.trim();
+      import.meta.hot.accept(${JSON.stringify(previewAnnotationURLs)}, (previewAnnotationModules) => {
+        // getProjectAnnotations has changed so we need to patch the new one in
+        window.__STORYBOOK_PREVIEW__.onGetProjectAnnotationsChanged({ getProjectAnnotations: () => getProjectAnnotations(previewAnnotationModules) });
+      });
+    }`.trim();
   };
 
   /**
@@ -63,18 +91,21 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
    * and the HMR implementation has been tweaked to work with Vite.
    * @todo Inline variable and remove `noinspection`
    */
-  const code = `
+  const code = dedent`
   import { composeConfigs, PreviewWeb, ClientApi } from 'storybook/internal/preview-api';
   import '${virtualAddonSetupFile}';
   import { importFn } from '${virtualStoriesFile}';
-
-    ${getPreviewAnnotationsFunction}
+  ${imports.join('\n')}
+  ${getPreviewAnnotationsFunction}
 
-    window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
-    
-    window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
-    
-    ${generateHMRHandler(frameworkName)};
-    `.trim();
+  window.__STORYBOOK_PREVIEW__ = window.__STORYBOOK_PREVIEW__ || new PreviewWeb(importFn, getProjectAnnotations);
+
+  window.__STORYBOOK_STORY_STORE__ = window.__STORYBOOK_STORY_STORE__ || window.__STORYBOOK_PREVIEW__.storyStore;
+
+  ${generateHMRHandler()};
+  `.trim();
   return code;
 }
+function hash(value: string) {
+  return value.split('').reduce((acc, char) => acc + char.charCodeAt(0), 0);
+}

code/builders/builder-vite/src/utils/process-preview-annotation.test.ts

@@ -3,64 +3,71 @@ import { processPreviewAnnotation } from './process-preview-annotation';
 import { onlyWindows, skipWindows } from '../../../../vitest.helpers';
 
 describe('processPreviewAnnotation()', () => {
-  it('should pull the `bare` value from an object', () => {
+  it('should pull the `absolute` value from an object', () => {
     const annotation = {
       bare: '@storybook/addon-links/preview',
       absolute: '/Users/foo/storybook/node_modules/@storybook/addon-links/dist/preview.mjs',
     };
     const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
-    expect(url).toBe('@storybook/addon-links/preview');
+    expect(url).toBe('/Users/foo/storybook/node_modules/@storybook/addon-links/dist/preview.mjs');
   });
 
-  it('should convert relative paths into urls', () => {
-    const annotation = './src/stories/components';
+  it('should convert relative paths into absolute paths', () => {
+    const annotation = './src/stories/preview';
     const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
-    expect(url).toBe('/src/stories/components');
+    expect(url).toBe('/Users/foo/storybook/src/stories/preview');
   });
 
-  skipWindows(() => {
-    it('should convert absolute filesystem paths into urls relative to project root', () => {
-      const annotation = '/Users/foo/storybook/.storybook/preview.js';
-      const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
-      expect(url).toBe('/.storybook/preview.js');
-    });
+  it('should keep absolute filesystem paths', () => {
+    const annotation = '/Users/foo/storybook/.storybook/preview.js';
+    const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
+    expect(url).toBe('/Users/foo/storybook/.storybook/preview.js');
+  });
 
-    // TODO: figure out why this fails on windows. Could be related to storybook-metadata.test file altering path.sep
-    it('should convert node_modules into bare paths', () => {
-      const annotation = '/Users/foo/storybook/node_modules/storybook-addon/preview';
-      const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
-      expect(url).toBe('storybook-addon/preview');
-    });
+  it('should keep absolute node_modules paths', () => {
+    const annotation = '/Users/foo/storybook/node_modules/storybook-addon/preview';
+    const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
+    expect(url).toBe('/Users/foo/storybook/node_modules/storybook-addon/preview');
+  });
 
-    it('should convert relative paths outside the root into absolute', () => {
-      const annotation = '../parent.js';
-      const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
-      expect(url).toBe('/Users/foo/parent.js');
-    });
+  it('should convert relative paths outside the root into absolute', () => {
+    const annotation = '../parent.js';
+    const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
+    expect(url).toBe('/Users/foo/parent.js');
+  });
 
-    it('should not change absolute paths outside of the project root', () => {
-      const annotation = '/Users/foo/parent.js';
-      const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
-      expect(url).toBe(annotation);
-    });
+  it('should not change absolute paths outside of the project root', () => {
+    const annotation = '/Users/foo/parent.js';
+    const url = processPreviewAnnotation(annotation, '/Users/foo/storybook/');
+    expect(url).toBe(annotation);
   });
 
-  onlyWindows(() => {
-    it('should convert absolute windows filesystem paths into urls relative to project root', () => {
-      const annotation = 'C:/foo/storybook/.storybook/preview.js';
-      const url = processPreviewAnnotation(annotation, 'C:/foo/storybook');
-      expect(url).toBe('/.storybook/preview.js');
-    });
-    it('should convert relative paths outside the root into absolute on Windows', () => {
-      const annotation = '../parent.js';
-      const url = processPreviewAnnotation(annotation, 'C:/Users/foo/storybook/');
-      expect(url).toBe('C:/Users/foo/parent.js');
-    });
+  it('should keep absolute windows filesystem paths as is', () => {
+    const annotation = 'C:/foo/storybook/.storybook/preview.js';
+    const url = processPreviewAnnotation(annotation, 'C:/foo/storybook');
+    expect(url).toBe('C:/foo/storybook/.storybook/preview.js');
+  });
+  it('should convert relative paths outside the root into absolute on Windows', () => {
+    const annotation = '../parent.js';
+    const url = processPreviewAnnotation(annotation, 'C:/Users/foo/storybook/');
+    expect(url).toBe('C:/Users/foo/parent.js');
+  });
+
+  it('should not change Windows absolute paths outside of the project root', () => {
+    const annotation = 'D:/Users/foo/parent.js';
+    const url = processPreviewAnnotation(annotation, 'D:/Users/foo/storybook/');
+    expect(url).toBe(annotation);
+  });
+
+  it('should normalize absolute Windows paths using \\', () => {
+    const annotation = 'C:\\foo\\storybook\\.storybook\\preview.js';
+    const url = processPreviewAnnotation(annotation, 'C:\\foo\\storybook');
+    expect(url).toBe('C:/foo/storybook/.storybook/preview.js');
+  });
 
-    it('should not change Windows absolute paths outside of the project root', () => {
-      const annotation = 'D:/Users/foo/parent.js';
-      const url = processPreviewAnnotation(annotation, 'D:/Users/foo/storybook/');
-      expect(url).toBe(annotation);
-    });
+  it('should normalize relative Windows paths using \\', () => {
+    const annotation = '.\\src\\stories\\preview';
+    const url = processPreviewAnnotation(annotation, 'C:\\foo\\storybook');
+    expect(url).toBe('C:/foo/storybook/src/stories/preview');
   });
 });