fix(addons): webcomponents transformSource

storybookjs · Jul 6, 2021 · f0bf666 · f0bf666
1 parent 65e7803
commit f0bf666
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/addons/docs/src/frameworks/web-components/sourceDecorator.test.ts b/addons/docs/src/frameworks/web-components/sourceDecorator.test.ts
@@ -11,7 +11,6 @@ expect.addSnapshotSerializer({
   test: (val) => typeof val === 'string',
 });
 
-// @ts-expect-error: :shrug:
 const makeContext = (name: string, parameters: any, args: any, extra?: object): StoryContext => ({
   id: `lit-test--${name}`,
   kind: 'js-text',
@@ -51,8 +50,8 @@ describe('sourceDecorator', () => {
   it('allows the snippet output to be modified by transformSource', () => {
     const storyFn = (args: any) => html`<div>args story</div>`;
     const transformSource = (dom: string) => `<p>${dom}</p>`;
-    const jsx = { transformSource };
-    const context = makeContext('args', { __isArgsStory: true, jsx }, {});
+    const docs = { transformSource };
+    const context = makeContext('args', { __isArgsStory: true, docs }, {});
     sourceDecorator(storyFn, context);
     expect(mockChannel.emit).toHaveBeenCalledWith(
       SNIPPET_RENDERED,
@@ -64,8 +63,8 @@ describe('sourceDecorator', () => {
   it('provides the story context to transformSource', () => {
     const storyFn = (args: any) => html`<div>args story</div>`;
     const transformSource = jest.fn();
-    const jsx = { transformSource };
-    const context = makeContext('args', { __isArgsStory: true, jsx }, {});
+    const docs = { transformSource };
+    const context = makeContext('args', { __isArgsStory: true, docs }, {});
     sourceDecorator(storyFn, context);
     expect(transformSource).toHaveBeenCalledWith('<div>args story</div>', context);
   });

diff --git a/addons/docs/src/frameworks/web-components/sourceDecorator.ts b/addons/docs/src/frameworks/web-components/sourceDecorator.ts
@@ -17,13 +17,19 @@ function skipSourceRender(context: StoryContext) {
   return !isArgsStory || sourceParams?.code || sourceParams?.type === SourceType.CODE;
 }
 
+function applyTransformSource(source, context) {
+  const { transformSource } = context.parameters.docs ?? {};
+  if (typeof transformSource !== 'function') return source;
+  return transformSource(source, context);
+}
+
 export function sourceDecorator(storyFn: StoryFn, context: StoryContext) {
   const story = storyFn();
 
   if (!skipSourceRender(context)) {
     const container = window.document.createElement('div');
     render(story, container);
-    const source = container.innerHTML.replace(/<!---->/g, '');
+    const source = applyTransformSource(container.innerHTML.replace(/<!---->/g, ''), context);
     if (source) addons.getChannel().emit(SNIPPET_RENDERED, context?.id, source);
   }