Send story context to transform source

addons/docs/docs/recipes.md

@@ -270,7 +270,7 @@ Alternatively, you can provide a function in the `docs.transformSource` paramete
 const SOURCE_REGEX = /^\(\) => `(.*)`$/;
 export const parameters = {
   docs: {
-    transformSource: (src, storyId) => {
+    transformSource: (src, storyId, storyContext) => {
       const match = SOURCE_REGEX.exec(src);
       return match ? match[1] : src;
     },

addons/docs/src/blocks/enhanceSource.test.ts

@@ -11,15 +11,14 @@ const emptyContext: StoryContext = {
   parameters: {},
 };
 
-const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
-
 describe('addon-docs enhanceSource', () => {
   describe('no source loaded', () => {
     const baseContext = emptyContext;
     it('no transformSource', () => {
       expect(enhanceSource(baseContext)).toBeNull();
     });
     it('transformSource', () => {
+      const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
       const parameters = { ...baseContext.parameters, docs: { transformSource } };
       expect(enhanceSource({ ...baseContext, parameters })).toBeNull();
     });
@@ -35,11 +34,23 @@ describe('addon-docs enhanceSource', () => {
       });
     });
     it('transformSource', () => {
+      const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
       const parameters = { ...baseContext.parameters, docs: { transformSource } };
       expect(enhanceSource({ ...baseContext, parameters }).docs.source).toEqual({
         code: 'formatted: storySource.source',
       });
     });
+    it('receives StoryContext as second argument', () => {
+      const transformSource = jest.fn();
+      const parameters = { ...baseContext.parameters, docs: { transformSource } };
+      const context = { ...baseContext, parameters };
+      enhanceSource(context);
+      expect(transformSource).toHaveBeenCalledWith(
+        baseContext.parameters.storySource.source,
+        baseContext.id,
+        context
+      );
+    });
   });
   describe('storysource source loaded w/ locationsMap', () => {
     const baseContext = {
@@ -57,11 +68,19 @@ describe('addon-docs enhanceSource', () => {
       expect(enhanceSource(baseContext)).toEqual({ docs: { source: { code: 'Source' } } });
     });
     it('transformSource', () => {
+      const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
       const parameters = { ...baseContext.parameters, docs: { transformSource } };
       expect(enhanceSource({ ...baseContext, parameters }).docs.source).toEqual({
         code: 'formatted: Source',
       });
     });
+    it('receives StoryContext as second argument', () => {
+      const transformSource = jest.fn();
+      const parameters = { ...baseContext.parameters, docs: { transformSource } };
+      const context = { ...baseContext, parameters };
+      enhanceSource(context);
+      expect(transformSource).toHaveBeenCalledWith('Source', context.id, context);
+    });
   });
   describe('custom docs.source provided', () => {
     const baseContext = {
@@ -75,9 +94,21 @@ describe('addon-docs enhanceSource', () => {
       expect(enhanceSource(baseContext)).toBeNull();
     });
     it('transformSource', () => {
+      const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
       const { source } = baseContext.parameters.docs;
       const parameters = { ...baseContext.parameters, docs: { source, transformSource } };
       expect(enhanceSource({ ...baseContext, parameters })).toBeNull();
     });
+    it('receives StoryContext as second argument', () => {
+      const transformSource = jest.fn();
+      const parameters = { ...baseContext.parameters, docs: { transformSource } };
+      const context = { ...baseContext, parameters };
+      enhanceSource(context);
+      expect(transformSource).toHaveBeenCalledWith(
+        context.parameters.storySource.source,
+        context.id,
+        context
+      );
+    });
   });
 });

addons/docs/src/blocks/enhanceSource.ts

@@ -54,7 +54,7 @@ export const enhanceSource = (context: StoryContext): Parameters => {
   }
 
   const input = extract(id, storySource);
-  const code = transformSource ? transformSource(input, id) : input;
+  const code = transformSource ? transformSource(input, id, context) : input;
 
   return { docs: combineParameters(docs, { source: { code } }) };
 };

examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot

@@ -2,7 +2,13 @@
 
 exports[`Storyshots Demo Button 1`] = `
 <button>
-  Hello Button
+  Click Me
+</button>
+`;
+
+exports[`Storyshots Demo Button With Custom Source Code 1`] = `
+<button>
+  Click Me
 </button>
 `;
 

examples/html-kitchen-sink/stories/button.stories.js

@@ -4,19 +4,42 @@ import { useEffect } from '@storybook/client-api';
 
 export default {
   title: 'Demo',
+  argTypes: {
+    buttonText: { control: 'text', defaultValue: 'Click Me' },
+  },
 };
 
-export const Heading = () => '<h1>Hello World</h1>';
-export const Headings = () =>
-  '<h1>Hello World</h1><h2>Hello World</h2><h3>Hello World</h3><h4>Hello World</h4>';
-
-export const Button = () => {
+const Template = ({ buttonText }) => {
   const btn = document.createElement('button');
-  btn.innerHTML = 'Hello Button';
+  btn.innerHTML = buttonText;
   btn.addEventListener('click', action('Click'));
   return btn;
 };
 
+export const ButtonWithCustomSourceCode = Template.bind({});
+
+ButtonWithCustomSourceCode.args = {
+  buttonText: 'Click Me',
+};
+
+ButtonWithCustomSourceCode.parameters = {
+  docs: {
+    transformSource: (src, id, context) => {
+      return Template(context.args).outerHTML;
+    },
+  },
+};
+
+export const Button = Template.bind({});
+
+Button.args = {
+  buttonText: 'Click Me',
+};
+
+export const Heading = () => '<h1>Hello World</h1>';
+export const Headings = () =>
+  '<h1>Hello World</h1><h2>Hello World</h2><h3>Hello World</h3><h4>Hello World</h4>';
+
 export const Effect = () => {
   useEffect(() => {
     document.getElementById('button').style.backgroundColor = 'yellow';

lib/api/src/version.ts

@@ -1 +1 @@
-export const version = '6.1.0-alpha.1';
+export const version = '6.1.0-alpha.1';

lib/cli/versions.json

@@ -52,4 +52,4 @@
   "@storybook/ui": "6.1.0-alpha.1",
   "@storybook/vue": "6.1.0-alpha.1",
   "@storybook/web-components": "6.1.0-alpha.1"
-}
+}