Merge pull request #16106 from storybookjs/16047-add-hmr-to-story-index

Refactor `stories-json` to use a caching class

lib/core-server/src/core-presets.test.ts

@@ -48,6 +48,7 @@ jest.mock('@storybook/store', () => {
   return {
     ...actualStore,
     autoTitle: () => 'auto-title',
+    autoTitleFromSpecifier: () => 'auto-title-from-specifier',
   };
 });
 

lib/core-server/src/utils/StoryIndexGenerator.ts

@@ -0,0 +1,144 @@
+import path from 'path';
+import fs from 'fs-extra';
+import glob from 'globby';
+
+import { autoTitleFromSpecifier, sortStories, Path, StoryIndex } from '@storybook/store';
+import { NormalizedStoriesSpecifier } from '@storybook/core-common';
+import { logger } from '@storybook/node-logger';
+import { readCsfOrMdx, getStorySortParameter } from '@storybook/csf-tools';
+
+function sortExtractedStories(
+  stories: StoryIndex['stories'],
+  storySortParameter: any,
+  fileNameOrder: string[]
+) {
+  const sortableStories = Object.entries(stories).map(([id, story]) => [
+    id,
+    { id, kind: story.title, story: story.name, ...story },
+    { fileName: story.importPath },
+  ]);
+  sortStories(sortableStories, storySortParameter, fileNameOrder);
+  return sortableStories.reduce((acc, item) => {
+    const storyId = item[0] as string;
+    acc[storyId] = stories[storyId];
+    return acc;
+  }, {} as StoryIndex['stories']);
+}
+
+export class StoryIndexGenerator {
+  // An internal cache mapping specifiers to a set of path=><set of stories>
+  // Later, we'll combine each of these subsets together to form the full index
+  private storyIndexEntries: Map<
+    NormalizedStoriesSpecifier,
+    Record<Path, StoryIndex['stories'] | false>
+  >;
+
+  constructor(
+    public readonly specifiers: NormalizedStoriesSpecifier[],
+    public readonly configDir: Path
+  ) {
+    this.storyIndexEntries = new Map();
+  }
+
+  async initialize() {
+    // Find all matching paths for each specifier
+    await Promise.all(
+      this.specifiers.map(async (specifier) => {
+        const pathToSubIndex = {} as Record<Path, StoryIndex['stories'] | false>;
+
+        const files = await glob(path.join(this.configDir, specifier.glob));
+        files.forEach((fileName: Path) => {
+          pathToSubIndex[fileName] = false;
+        });
+
+        this.storyIndexEntries.set(specifier, pathToSubIndex);
+      })
+    );
+
+    // Extract stories for each file
+    await this.ensureExtracted();
+  }
+
+  async ensureExtracted() {
+    await Promise.all(
+      this.specifiers.map(async (specifier) => {
+        const entry = this.storyIndexEntries.get(specifier);
+        await Promise.all(
+          Object.keys(entry).map(async (fileName) => {
+            if (!entry[fileName]) await this.extractStories(specifier, fileName);
+          })
+        );
+      })
+    );
+  }
+
+  async extractStories(specifier: NormalizedStoriesSpecifier, absolutePath: Path) {
+    const ext = path.extname(absolutePath);
+    const relativePath = path.relative(this.configDir, absolutePath);
+    if (!['.js', '.jsx', '.ts', '.tsx', '.mdx'].includes(ext)) {
+      logger.info(`Skipping ${ext} file ${relativePath}`);
+      return;
+    }
+    try {
+      const entry = this.storyIndexEntries.get(specifier);
+      const fileStories = {} as StoryIndex['stories'];
+
+      const importPath = relativePath[0] === '.' ? relativePath : `./${relativePath}`;
+      const defaultTitle = autoTitleFromSpecifier(importPath, specifier);
+      const csf = (await readCsfOrMdx(absolutePath, { defaultTitle })).parse();
+      csf.stories.forEach(({ id, name }) => {
+        fileStories[id] = {
+          title: csf.meta.title,
+          name,
+          importPath,
+        };
+      });
+
+      entry[absolutePath] = fileStories;
+    } catch (err) {
+      logger.warn(`🚨 Extraction error on ${relativePath}: ${err}`);
+      logger.warn(`🚨 ${err.stack}`);
+      throw err;
+    }
+  }
+
+  async getIndex() {
+    // Extract any entries that are currently missing
+    await this.ensureExtracted();
+
+    const stories: StoryIndex['stories'] = {};
+
+    // Check each entry and compose into stories, extracting if needed
+    this.specifiers.map(async (specifier) => {
+      Object.values(this.storyIndexEntries.get(specifier)).map((subStories) =>
+        Object.assign(stories, subStories)
+      );
+    });
+
+    const storySortParameter = await this.getStorySortParameter();
+    const sorted = sortExtractedStories(stories, storySortParameter, this.storyFileNames());
+
+    return {
+      v: 3,
+      stories: sorted,
+    };
+  }
+
+  async getStorySortParameter() {
+    const previewFile = ['js', 'jsx', 'ts', 'tsx']
+      .map((ext) => path.join(this.configDir, `preview.${ext}`))
+      .find((fname) => fs.existsSync(fname));
+    let storySortParameter;
+    if (previewFile) {
+      const previewCode = (await fs.readFile(previewFile, 'utf-8')).toString();
+      storySortParameter = await getStorySortParameter(previewCode);
+    }
+
+    return storySortParameter;
+  }
+
+  // Get the story file names in "imported order"
+  storyFileNames() {
+    return Array.from(this.storyIndexEntries.values()).flatMap((r) => Object.keys(r));
+  }
+}

lib/core-server/src/utils/stories-json.ts

@@ -1,98 +1,17 @@
-import path from 'path';
 import fs from 'fs-extra';
-import glob from 'globby';
-import { logger } from '@storybook/node-logger';
 import { Options, normalizeStories, NormalizedStoriesSpecifier } from '@storybook/core-common';
-import { autoTitle, sortStories } from '@storybook/store';
-import { readCsfOrMdx, getStorySortParameter } from '@storybook/csf-tools';
-
-interface ExtractedStory {
-  title: string;
-  name: string;
-  importPath: string;
-}
-
-type ExtractedStories = Record<string, ExtractedStory>;
-
-function sortExtractedStories(
-  stories: ExtractedStories,
-  storySortParameter: any,
-  fileNameOrder: string[]
-) {
-  const sortableStories = Object.entries(stories).map(([id, story]) => [
-    id,
-    { id, kind: story.title, story: story.name, ...story },
-    { fileName: story.importPath },
-  ]);
-  sortStories(sortableStories, storySortParameter, fileNameOrder);
-  return sortableStories.reduce((acc, item) => {
-    const storyId = item[0] as string;
-    acc[storyId] = stories[storyId];
-    return acc;
-  }, {} as ExtractedStories);
-}
-
-async function extractStories(normalizedStories: NormalizedStoriesSpecifier[], configDir: string) {
-  const storiesGlobs = normalizedStories.map((s) => s.glob);
-  const storyFiles: string[] = [];
-  await Promise.all(
-    storiesGlobs.map(async (storiesGlob) => {
-      const files = await glob(path.join(configDir, storiesGlob));
-      storyFiles.push(...files);
-    })
-  );
-  logger.info(`⚙️ Processing ${storyFiles.length} story files from ${storiesGlobs}`);
-
-  const stories: ExtractedStories = {};
-  await Promise.all(
-    storyFiles.map(async (absolutePath) => {
-      const ext = path.extname(absolutePath);
-      const relativePath = path.relative(configDir, absolutePath);
-      if (!['.js', '.jsx', '.ts', '.tsx', '.mdx'].includes(ext)) {
-        logger.info(`Skipping ${ext} file ${relativePath}`);
-        return;
-      }
-      try {
-        const importPath = relativePath[0] === '.' ? relativePath : `./${relativePath}`;
-        const defaultTitle = autoTitle(importPath, normalizedStories);
-        const csf = (await readCsfOrMdx(absolutePath, { defaultTitle })).parse();
-        csf.stories.forEach(({ id, name }) => {
-          stories[id] = {
-            title: csf.meta.title,
-            name,
-            importPath,
-          };
-        });
-      } catch (err) {
-        logger.warn(`🚨 Extraction error on ${relativePath}: ${err}`);
-        logger.warn(`🚨 ${err.stack}`);
-        throw err;
-      }
-    })
-  );
-
-  const previewFile = ['js', 'jsx', 'ts', 'tsx']
-    .map((ext) => path.join(configDir, `preview.${ext}`))
-    .find((fname) => fs.existsSync(fname));
-
-  let storySortParameter;
-  if (previewFile) {
-    const previewCode = (await fs.readFile(previewFile, 'utf-8')).toString();
-    storySortParameter = await getStorySortParameter(previewCode);
-  }
-
-  const sorted = sortExtractedStories(stories, storySortParameter, storyFiles);
-
-  return sorted;
-}
+import { StoryIndexGenerator } from './StoryIndexGenerator';
 
 export async function extractStoriesJson(
   outputFile: string,
   normalizedStories: NormalizedStoriesSpecifier[],
   configDir: string
 ) {
-  const stories = await extractStories(normalizedStories, configDir);
-  await fs.writeJson(outputFile, { v: 3, stories });
+  const generator = new StoryIndexGenerator(normalizedStories, configDir);
+  await generator.initialize();
+
+  const index = await generator.getIndex();
+  await fs.writeJson(outputFile, index);
 }
 
 export async function useStoriesJson(router: any, options: Options) {
@@ -100,19 +19,17 @@ export async function useStoriesJson(router: any, options: Options) {
     configDir: options.configDir,
     workingDir: process.cwd(),
   });
+
   router.use('/stories.json', async (_req: any, res: any) => {
-    extractStories(normalized, options.configDir)
-      .then((stories: ExtractedStories) => {
-        res.header('Content-Type', 'application/json');
-        return res.send(
-          JSON.stringify({
-            v: 3,
-            stories,
-          })
-        );
-      })
-      .catch((err: Error) => {
-        res.status(500).send(err.message);
-      });
+    const generator = new StoryIndexGenerator(normalized, options.configDir);
+    await generator.initialize();
+
+    try {
+      const index = await generator.getIndex();
+      res.header('Content-Type', 'application/json');
+      return res.send(JSON.stringify(index));
+    } catch (err) {
+      return res.status(500).send(err.message);
+    }
   });
 }

lib/store/src/autoTitle.test.ts

@@ -1,4 +1,4 @@
-import { autoTitleFromEntry as auto } from './autoTitle';
+import { autoTitleFromSpecifier as auto } from './autoTitle';
 
 expect.addSnapshotSerializer({
   print: (val: any) => val,

lib/store/src/autoTitle.ts

@@ -32,7 +32,7 @@ const startCaseTitle = (title: string) => {
   return title.split('/').map(startCase).join('/');
 };
 
-export const autoTitleFromEntry = (fileName: string, entry: NormalizedStoriesSpecifier) => {
+export const autoTitleFromSpecifier = (fileName: string, entry: NormalizedStoriesSpecifier) => {
   const { directory, titlePrefix = '' } = entry.specifier || {};
   // On Windows, backslashes are used in paths, which can cause problems here
   // slash makes sure we always handle paths with unix-style forward slash
@@ -49,7 +49,7 @@ export const autoTitleFromEntry = (fileName: string, entry: NormalizedStoriesSpe
 
 export const autoTitle = (fileName: string, storiesEntries: NormalizedStoriesSpecifier[]) => {
   for (let i = 0; i < storiesEntries.length; i += 1) {
-    const title = autoTitleFromEntry(fileName, storiesEntries[i]);
+    const title = autoTitleFromSpecifier(fileName, storiesEntries[i]);
     if (title) return title;
   }
   return undefined;