Docs: Support async docs.source.transform

code/addons/docs/docs/recipes.md

@@ -260,7 +260,7 @@ Example.parameters = {
 };
 ```
 
-Alternatively, you can provide a function in the `docs.source.transform` parameter. For example, the following snippet in `.storybook/preview.js` globally removes the arrow at the beginning of a function that returns a string:
+Alternatively, you can provide a function or an async function in the `docs.source.transform` parameter. For example, the following snippet in `.storybook/preview.js` globally removes the arrow at the beginning of a function that returns a string:
 
 ```js
 const SOURCE_REGEX = /^\(\) => `(.*)`$/;

code/addons/docs/template/stories/docspage/source.stories.ts

@@ -49,3 +49,23 @@ export const Transform = {
     },
   },
 };
+
+export const AsyncTransform = {
+  parameters: {
+    docs: {
+      source: {
+        async transform(src: string, storyContext: StoryContext) {
+          await Promise.resolve<any>((res) =>
+            setTimeout(() => {
+              res();
+            }, 500)
+          );
+          return dedent`// We transformed this asynchronously!
+          // The current args are: ${JSON.stringify(storyContext.args)}
+          const example = (${src});
+          `;
+        },
+      },
+    },
+  },
+};

code/lib/blocks/src/blocks/Source.tsx

@@ -1,8 +1,8 @@
 import type { ComponentProps, FC } from 'react';
-import React, { useContext } from 'react';
+import React, { useContext, useEffect, useMemo, useState } from 'react';
 
 import { SourceType } from 'storybook/internal/docs-tools';
-import type { Args, ModuleExport, PreparedStory, StoryId } from 'storybook/internal/types';
+import type { Args, ModuleExport, StoryId } from 'storybook/internal/types';
 
 import type { SourceCodeProps } from '../components/Source';
 import { Source as PureSource, SourceError } from '../components/Source';
@@ -18,7 +18,7 @@ type SourceParameters = SourceCodeProps & {
   transform?: (
     code: string,
     storyContext: ReturnType<DocsContextProps['getStoryContext']>
-  ) => string;
+  ) => string | Promise<string>;
   /** Internal: set by our CSF loader (`enrichCsf` in `@storybook/csf-tools`). */
   originalSource?: string;
 };
@@ -58,7 +58,7 @@ const getStorySource = (
   return source || { code: '' };
 };
 
-const getSnippet = ({
+const useCode = ({
   snippet,
   storyContext,
   typeFromProps,
@@ -70,15 +70,11 @@ const getSnippet = ({
   transformFromProps?: SourceProps['transform'];
 }): string => {
   const { __isArgsStory: isArgsStory } = storyContext.parameters;
+  const [transformedCode, setTransformedCode] = useState('Loading...');
   const sourceParameters = (storyContext.parameters.docs?.source || {}) as SourceParameters;
 
   const type = typeFromProps || sourceParameters.type || SourceType.AUTO;
 
-  // if user has hard-coded the snippet, that takes precedence
-  if (sourceParameters.code !== undefined) {
-    return sourceParameters.code;
-  }
-
   const useSnippet =
     // if user has explicitly set this as dynamic, use snippet
     type === SourceType.DYNAMIC ||
@@ -88,8 +84,29 @@ const getSnippet = ({
   const code = useSnippet ? snippet : sourceParameters.originalSource || '';
 
   const transformer = transformFromProps ?? sourceParameters.transform;
+  const transformed = useMemo(
+    () => (transformer ? transformer?.(code, storyContext) : code),
+    [code, storyContext, transformer]
+  );
+
+  useEffect(() => {
+    async function getTransformedCode() {
+      setTransformedCode(await transformed);
+    }
 
-  return transformer?.(code, storyContext) || code;
+    getTransformedCode();
+  }, [transformed]);
+
+  if (sourceParameters.code !== undefined) {
+    return sourceParameters.code;
+  }
+
+  if (typeof transformed === 'object' && typeof transformed.then === 'function') {
+    console.log({ transformedCode });
+    return transformedCode;
+  }
+
+  return transformed as string;
 };
 
 // state is used by the Canvas block, which also calls useSourceProps
@@ -100,58 +117,64 @@ export const useSourceProps = (
   docsContext: DocsContextProps<any>,
   sourceContext: SourceContextProps
 ): PureSourceProps => {
-  let story: PreparedStory;
   const { of } = props;
-  if ('of' in props && of === undefined) {
-    throw new Error('Unexpected `of={undefined}`, did you mistype a CSF file reference?');
-  }
 
-  if (of) {
-    const resolved = docsContext.resolveOf(of, ['story']);
-    story = resolved.story;
-  } else {
-    try {
-      // Always fall back to the primary story for source parameters, even if code is set.
-      story = docsContext.storyById();
-    } catch (err) {
-      // You are allowed to use <Source code="..." /> and <Canvas /> unattached.
+  const story = useMemo(() => {
+    if (of) {
+      const resolved = docsContext.resolveOf(of, ['story']);
+      return resolved.story;
+    } else {
+      try {
+        // Always fall back to the primary story for source parameters, even if code is set.
+        return docsContext.storyById();
+      } catch (err) {
+        // You are allowed to use <Source code="..." /> and <Canvas /> unattached.
+      }
     }
+  }, [docsContext, of]);
+
+  const storyContext = docsContext.getStoryContext(story);
+
+  // eslint-disable-next-line no-underscore-dangle
+  const argsForSource = props.__forceInitialArgs
+    ? storyContext.initialArgs
+    : storyContext.unmappedArgs;
+
+  const source = story ? getStorySource(story.id, argsForSource, sourceContext) : null;
+
+  const transformedCode = useCode({
+    snippet: source ? source.code : '',
+    storyContext: { ...storyContext, args: argsForSource },
+    typeFromProps: props.type,
+    transformFromProps: props.transform,
+  });
+
+  if ('of' in props && of === undefined) {
+    throw new Error('Unexpected `of={undefined}`, did you mistype a CSF file reference?');
   }
 
   const sourceParameters = (story?.parameters?.docs?.source || {}) as SourceParameters;
-  const { code } = props; // We will fall back to `sourceParameters.code`, but per story below
   let format = props.format ?? sourceParameters.format;
   const language = props.language ?? sourceParameters.language ?? 'jsx';
   const dark = props.dark ?? sourceParameters.dark ?? false;
 
-  if (!code && !story) {
+  if (!props.code && !story) {
     return { error: SourceError.SOURCE_UNAVAILABLE };
   }
-  if (code) {
+
+  if (props.code) {
     return {
-      code,
+      code: props.code,
       format,
       language,
       dark,
     };
   }
-  const storyContext = docsContext.getStoryContext(story);
-
-  // eslint-disable-next-line no-underscore-dangle
-  const argsForSource = props.__forceInitialArgs
-    ? storyContext.initialArgs
-    : storyContext.unmappedArgs;
 
-  const source = getStorySource(story.id, argsForSource, sourceContext);
   format = source.format ?? story.parameters.docs?.source?.format ?? false;
 
   return {
-    code: getSnippet({
-      snippet: source.code,
-      storyContext: { ...storyContext, args: argsForSource },
-      typeFromProps: props.type,
-      transformFromProps: props.transform,
-    }),
+    code: transformedCode,
     format,
     language,
     dark,

docs/api/doc-blocks/doc-block-source.mdx

@@ -102,12 +102,14 @@ Determines if the snippet is rendered in dark mode.
   Determines if [decorators](../../writing-stories/decorators.mdx) are rendered in the source code snippet.
 </IfRenderer>
 
-### `format`
+### `format` (deprecated - will be removed in Storybook 9)
 
 Type: `boolean | 'dedent' | BuiltInParserName`
 
 Default: `parameters.docs.source.format` or `true`
 
+Deprecated: Will be removed in Storybook 9. Please use `docs.source.transform` instead and use your own formatter
+
 Specifies the formatting used on source code. Both `true` and `'dedent'` have the same effect of removing any extraneous indentation. Supports all valid [prettier parser names](https://prettier.io/docs/en/configuration.html#setting-the-parserdocsenoptionshtmlparser-option).
 
 ### `language`
@@ -138,7 +140,7 @@ Type: `(code: string, storyContext: StoryContext) => string`
 
 Default: `parameters.docs.source.transform`
 
-A function to dynamically transform the source before being rendered, based on the original source and any story context necessary. The returned string is displayed as-is.
+An async function to dynamically transform the source before being rendered, based on the original source and any story context necessary. The returned string is displayed as-is.
 If both [`code`](#code) and `transform` are specified, `transform` will be ignored.
 
 ### `type`