Preview API: Add csf factory utilities

code/addons/test/src/vitest-plugin/test-utils.ts

@@ -3,7 +3,11 @@
 /* eslint-disable no-underscore-dangle */
 import { type RunnerTask, type TaskMeta, type TestContext } from 'vitest';
 
-import { type Report, composeStory } from 'storybook/internal/preview-api';
+import {
+  type Report,
+  composeStory,
+  getCsfFactoryAnnotations,
+} from 'storybook/internal/preview-api';
 import type { ComponentAnnotations, ComposedStoryFn } from 'storybook/internal/types';
 
 import { server } from '@vitest/browser/context';
@@ -26,11 +30,12 @@ export const testStory = (
   skipTags: string[]
 ) => {
   return async (context: TestContext & { story: ComposedStoryFn }) => {
+    const annotations = getCsfFactoryAnnotations(story, meta);
     const composedStory = composeStory(
-      'isCSFFactory' in story ? (story as any).input : story,
-      'isCSFFactory' in story ? (meta as any).input : meta,
+      annotations.story,
+      annotations.meta,
       { initialGlobals: (await getInitialGlobals?.()) ?? {}, tags: await getTags?.() },
-      'isCSFFactory' in story ? (story as any).config?.input : undefined,
+      annotations.preview,
       exportName
     );
 

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

@@ -24,9 +24,7 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
   const getPreviewAnnotationsFunction = `
   const getProjectAnnotations = async (hmrPreviewAnnotationModules = []) => {
     const preview = await import('${previewFileUrl}');
-    const csfFactoryPreview = Object.values(preview).find(module => {
-      return 'isCSFFactoryPreview' in module
-    });
+    const csfFactoryPreview = getCsfFactoryPreview(preview);
 
     if (csfFactoryPreview) {
       return csfFactoryPreview.input;
@@ -81,7 +79,7 @@ export async function generateModernIframeScriptCode(options: Options, projectRo
 
   setup();
 
-  import { composeConfigs, PreviewWeb, ClientApi } from 'storybook/internal/preview-api';
+  import { composeConfigs, PreviewWeb, ClientApi, getCsfFactoryPreview } from 'storybook/internal/preview-api';
   import { importFn } from '${SB_VIRTUAL_FILES.VIRTUAL_STORIES_FILE}';
 
 

code/core/src/preview-api/index.ts

@@ -56,7 +56,13 @@ export {
   normalizeProjectAnnotations,
 } from './store';
 
-export { createPlaywrightTest } from './modules/store/csf/portable-stories';
+/** CSF API */
+export {
+  createPlaywrightTest,
+  getCsfFactoryPreview,
+  getCsfFactoryAnnotations,
+  isCsfFactory,
+} from './modules/store/csf';
 
 export type { PropDescriptor } from './store';
 

code/core/src/preview-api/modules/preview-web/docs-context/DocsContext.ts

@@ -13,7 +13,7 @@ import type {
 
 import { dedent } from 'ts-dedent';
 
-import type { StoryStore } from '../../store';
+import { type StoryStore, isCsfFactory } from '../../store';
 import type { DocsContextProps } from './DocsContextProps';
 
 export class DocsContext<TRenderer extends Renderer> implements DocsContextProps<TRenderer> {
@@ -163,7 +163,8 @@ export class DocsContext<TRenderer extends Renderer> implements DocsContextProps
     }
 
     const story = this.exportToStory.get(
-      'isCSFFactory' in moduleExportOrType ? moduleExportOrType.input : moduleExportOrType
+      // TODO: @kasperpeulen will fix this once csf factory types are defined
+      isCsfFactory(moduleExportOrType) ? (moduleExportOrType as any).input : moduleExportOrType
     );
 
     if (story) {

code/core/src/preview-api/modules/store/csf/csf-factory-utils.ts

@@ -0,0 +1,42 @@
+/* eslint-disable @typescript-eslint/naming-convention */
+
+/* eslint-disable no-underscore-dangle */
+import type {
+  Args,
+  ComponentAnnotations,
+  LegacyStoryAnnotationsOrFn,
+  ModuleExports,
+  ProjectAnnotations,
+  Renderer,
+  StoryAnnotations,
+} from '@storybook/types';
+
+export function getCsfFactoryPreview(preview: ModuleExports): ProjectAnnotations<any> | null {
+  return Object.values(preview).find(isCsfFactory) ?? null;
+}
+
+export function isCsfFactory(target: StoryAnnotations | ProjectAnnotations<any>) {
+  return (
+    target != null &&
+    typeof target === 'object' &&
+    ('isCSFFactory' in target || 'isCSFFactoryPreview' in target)
+  );
+}
+
+export function getCsfFactoryAnnotations<
+  TRenderer extends Renderer = Renderer,
+  TArgs extends Args = Args,
+>(
+  story: LegacyStoryAnnotationsOrFn<TRenderer>,
+  meta?: ComponentAnnotations<TRenderer, TArgs>,
+  projectAnnotations?: ProjectAnnotations<TRenderer>
+) {
+  const _isCsfFactory = isCsfFactory(story);
+
+  return {
+    // TODO: @kasperpeulen will fix this once csf factory types are defined
+    story: _isCsfFactory ? (story as any)?.input : story,
+    meta: _isCsfFactory ? (story as any)?.meta?.input : meta,
+    preview: _isCsfFactory ? (story as any)?.config?.input : projectAnnotations,
+  };
+}

code/core/src/preview-api/modules/store/csf/index.ts

@@ -8,3 +8,4 @@ export * from './getValuesFromArgTypes';
 export * from './composeConfigs';
 export * from './stepRunners';
 export * from './portable-stories';
+export * from './csf-factory-utils';

code/core/src/preview-api/modules/store/csf/portable-stories.ts

@@ -28,6 +28,7 @@ import { dedent } from 'ts-dedent';
 import { HooksContext } from '../../../addons';
 import { ReporterAPI } from '../reporter-api';
 import { composeConfigs } from './composeConfigs';
+import { getCsfFactoryAnnotations } from './csf-factory-utils';
 import { getValuesFromArgTypes } from './getValuesFromArgTypes';
 import { normalizeComponentAnnotations } from './normalizeComponentAnnotations';
 import { normalizeProjectAnnotations } from './normalizeProjectAnnotations';
@@ -89,21 +90,6 @@ export function setProjectAnnotations<TRenderer extends Renderer = Renderer>(
 
 const cleanups: CleanupCallback[] = [];
 
-export function getAnnotations<TRenderer extends Renderer = Renderer, TArgs extends Args = Args>(
-  story: LegacyStoryAnnotationsOrFn<TRenderer>,
-  meta?: ComponentAnnotations<TRenderer, TArgs>,
-  projectAnnotations?: ProjectAnnotations<TRenderer>
-) {
-  const isCsfFactory =
-    (typeof story === 'function' || typeof story === 'object') && 'isCSFFactory' in story;
-
-  return {
-    storyAnnotations: isCsfFactory ? (story as any)?.input : story,
-    componentAnnotations: isCsfFactory ? (story as any)?.meta?.input : meta,
-    projectAnnotations: isCsfFactory ? (story as any)?.config?.input : projectAnnotations,
-  };
-}
-
 export function composeStory<TRenderer extends Renderer = Renderer, TArgs extends Args = Args>(
   storyAnnotations: LegacyStoryAnnotationsOrFn<TRenderer>,
   componentAnnotations: ComponentAnnotations<TRenderer, TArgs>,
@@ -295,7 +281,8 @@ export function composeStories<TModule extends Store_CSFExports>(
 
   const composedStories = Object.entries(stories).reduce(
     (storiesMap, [exportsName, story]: [string, any]) => {
-      const { storyAnnotations, componentAnnotations } = getAnnotations(story);
+      const { story: storyAnnotations, meta: componentAnnotations } =
+        getCsfFactoryAnnotations(story);
       if (!meta && componentAnnotations) {
         meta = componentAnnotations;
       }

code/core/src/preview-api/modules/store/csf/processCSFFile.ts

@@ -48,6 +48,7 @@ export function processCSFFile<TRenderer extends Renderer>(
 
   const firstStory: any = Object.values(namedExports)[0];
   // CSF4
+  // TODO: @kasperpeulen will fix this once csf factory types are defined
   if (!defaultExport && 'isCSFFactory' in firstStory) {
     const meta: NormalizedComponentAnnotations<TRenderer> =
       normalizeComponentAnnotations<TRenderer>(firstStory.meta.input, title, importPath);