import memoize from 'memoizerific'; import type { Parameters, StoryId, StoryContextForLoaders, AnyFramework, ProjectAnnotations, ComponentTitle, StoryContextForEnhancers, StoryContext, } from '@storybook/csf'; import mapValues from 'lodash/mapValues'; import pick from 'lodash/pick'; import global from 'global'; import { SynchronousPromise } from 'synchronous-promise'; import { StoryIndexStore } from './StoryIndexStore'; import { ArgsStore } from './ArgsStore'; import { GlobalsStore } from './GlobalsStore'; import { processCSFFile, prepareStory, normalizeProjectAnnotations } from './csf'; import type { CSFFile, ModuleImportFn, Story, NormalizedProjectAnnotations, Path, ExtractOptions, BoundStory, PromiseLike, StoryIndex, StoryIndexEntry, V2CompatIndexEntry, } from './types'; import { HooksContext } from './hooks'; // TODO -- what are reasonable values for these? const CSF_CACHE_SIZE = 1000; const STORY_CACHE_SIZE = 10000; export class StoryStore { storyIndex: StoryIndexStore; importFn: ModuleImportFn; projectAnnotations: NormalizedProjectAnnotations; globals: GlobalsStore; args: ArgsStore; hooks: Record>; cachedCSFFiles?: Record>; processCSFFileWithCache: typeof processCSFFile; prepareStoryWithCache: typeof prepareStory; initializationPromise: SynchronousPromise; resolveInitializationPromise: () => void; constructor() { this.globals = new GlobalsStore(); this.args = new ArgsStore(); this.hooks = {}; // We use a cache for these two functions for two reasons: // 1. For performance // 2. To ensure that when the same story is prepared with the same inputs you get the same output this.processCSFFileWithCache = memoize(CSF_CACHE_SIZE)(processCSFFile) as typeof processCSFFile; this.prepareStoryWithCache = memoize(STORY_CACHE_SIZE)(prepareStory) as typeof prepareStory; // We cannot call `loadStory()` until we've been initialized properly. But we can wait for it. this.initializationPromise = new SynchronousPromise((resolve) => { this.resolveInitializationPromise = resolve; }); } setProjectAnnotations(projectAnnotations: ProjectAnnotations) { // By changing `this.projectAnnotations, we implicitly invalidate the `prepareStoryWithCache` this.projectAnnotations = normalizeProjectAnnotations(projectAnnotations); const { globals, globalTypes } = projectAnnotations; this.globals.set({ globals, globalTypes }); } initialize({ storyIndex, importFn, cache = false, }: { storyIndex?: StoryIndex; importFn: ModuleImportFn; cache?: boolean; }): PromiseLike { this.storyIndex = new StoryIndexStore(storyIndex); this.importFn = importFn; // We don't need the cache to be loaded to call `loadStory`, we just need the index ready this.resolveInitializationPromise(); return cache ? this.cacheAllCSFFiles() : SynchronousPromise.resolve(); } // This means that one of the CSF files has changed. // If the `importFn` has changed, we will invalidate both caches. // If the `storyIndex` data has changed, we may or may not invalidate the caches, depending // on whether we've loaded the relevant files yet. async onStoriesChanged({ importFn, storyIndex, }: { importFn?: ModuleImportFn; storyIndex?: StoryIndex; }) { if (importFn) this.importFn = importFn; if (storyIndex) this.storyIndex.stories = storyIndex.stories; if (this.cachedCSFFiles) await this.cacheAllCSFFiles(); } // To load a single CSF file to service a story we need to look up the importPath in the index loadCSFFileByStoryId(storyId: StoryId): PromiseLike> { const { importPath, title } = this.storyIndex.storyIdToEntry(storyId); return this.importFn(importPath).then((moduleExports) => // We pass the title in here as it may have been generated by autoTitle on the server. this.processCSFFileWithCache(moduleExports, importPath, title) ); } loadAllCSFFiles(): PromiseLike['cachedCSFFiles']> { const importPaths: Record = {}; Object.entries(this.storyIndex.stories).forEach(([storyId, { importPath }]) => { importPaths[importPath] = storyId; }); const csfFilePromiseList = Object.entries(importPaths).map(([importPath, storyId]) => this.loadCSFFileByStoryId(storyId).then((csfFile) => ({ importPath, csfFile, })) ); return SynchronousPromise.all(csfFilePromiseList).then((list) => list.reduce((acc, { importPath, csfFile }) => { acc[importPath] = csfFile; return acc; }, {} as Record>) ); } cacheAllCSFFiles(): PromiseLike { return this.initializationPromise.then(() => this.loadAllCSFFiles().then((csfFiles) => { this.cachedCSFFiles = csfFiles; }) ); } // Load the CSF file for a story and prepare the story from it and the project annotations. async loadStory({ storyId }: { storyId: StoryId }): Promise> { await this.initializationPromise; const csfFile = await this.loadCSFFileByStoryId(storyId); return this.storyFromCSFFile({ storyId, csfFile }); } // This function is synchronous for convenience -- often times if you have a CSF file already // it is easier not to have to await `loadStory`. storyFromCSFFile({ storyId, csfFile, }: { storyId: StoryId; csfFile: CSFFile; }): Story { const storyAnnotations = csfFile.stories[storyId]; if (!storyAnnotations) { throw new Error(`Didn't find '${storyId}' in CSF file, this is unexpected`); } const componentAnnotations = csfFile.meta; const story = this.prepareStoryWithCache( storyAnnotations, componentAnnotations, this.projectAnnotations ); this.args.setInitial(story); this.hooks[story.id] = this.hooks[story.id] || new HooksContext(); return story; } // If we have a CSF file we can get all the stories from it synchronously componentStoriesFromCSFFile({ csfFile }: { csfFile: CSFFile }): Story[] { return Object.keys(this.storyIndex.stories) .filter((storyId: StoryId) => !!csfFile.stories[storyId]) .map((storyId: StoryId) => this.storyFromCSFFile({ storyId, csfFile })); } // A prepared story does not include args, globals or hooks. These are stored in the story store // and updated separtely to the (immutable) story. getStoryContext(story: Story): Omit, 'viewMode'> { return { ...story, args: this.args.get(story.id), globals: this.globals.get(), hooks: this.hooks[story.id] as unknown, }; } cleanupStory(story: Story): void { this.hooks[story.id].clean(); } extract( options: ExtractOptions = { includeDocsOnly: false } ): Record> { if (!this.cachedCSFFiles) { throw new Error('Cannot call extract() unless you call cacheAllCSFFiles() first.'); } return Object.entries(this.storyIndex.stories).reduce((acc, [storyId, { importPath }]) => { const csfFile = this.cachedCSFFiles[importPath]; const story = this.storyFromCSFFile({ storyId, csfFile }); if (!options.includeDocsOnly && story.parameters.docsOnly) { return acc; } acc[storyId] = Object.entries(story).reduce( (storyAcc, [key, value]) => { if (typeof value === 'function') { return storyAcc; } if (Array.isArray(value)) { return Object.assign(storyAcc, { [key]: value.slice().sort() }); } return Object.assign(storyAcc, { [key]: value }); }, { args: story.initialArgs } ); return acc; }, {} as Record); } getSetStoriesPayload() { const stories = this.extract({ includeDocsOnly: true }); const kindParameters: Parameters = Object.values(stories).reduce( (acc: Parameters, { title }: { title: ComponentTitle }) => { acc[title] = {}; return acc; }, {} as Parameters ); return { v: 2, globals: this.globals.get(), globalParameters: {}, kindParameters, stories, }; } getStoriesJsonData = () => { const value = this.getSetStoriesPayload(); const allowedParameters = ['fileName', 'docsOnly', 'framework', '__id', '__isArgsStory']; const stories: Record = mapValues( value.stories, (story) => ({ ...pick(story, ['id', 'name', 'title']), importPath: this.storyIndex.stories[story.id].importPath, ...(!global.FEATURES?.breakingChangesV7 && { kind: story.title, story: story.name, parameters: { ...pick(story.parameters, allowedParameters), fileName: this.storyIndex.stories[story.id].importPath, }, }), }) ); return { v: 3, stories, }; }; raw(): BoundStory[] { return Object.values(this.extract()).map(({ id }: { id: StoryId }) => this.fromId(id)); } fromId(storyId: StoryId): BoundStory { if (!this.cachedCSFFiles) { throw new Error('Cannot call fromId/raw() unless you call cacheAllCSFFiles() first.'); } let importPath; try { ({ importPath } = this.storyIndex.storyIdToEntry(storyId)); } catch (err) { return null; } const csfFile = this.cachedCSFFiles[importPath]; const story = this.storyFromCSFFile({ storyId, csfFile }); return { ...story, storyFn: (update) => { const context = { ...this.getStoryContext(story), viewMode: 'story', } as StoryContext; return story.unboundStoryFn({ ...context, ...update }); }, }; } }