Ensure we emit all SNIPPET_RENDERED in useEffect()

This also enables us to simplify `SourceContainer`.

addons/docs/src/blocks/SourceContainer.tsx

@@ -18,24 +18,18 @@ export const SourceContainer: FC<{}> = ({ children }) => {
   const [sources, setSources] = useState<StorySources>({});
   const channel = addons.getChannel();
 
-  const sourcesRef = React.useRef<StorySources>();
-  const handleSnippetRendered = (id: StoryId, newItem: SourceItem) => {
-    if (newItem !== sources[id]) {
-      const newSources = { ...sourcesRef.current, [id]: newItem };
-      sourcesRef.current = newSources;
-    }
-  };
-
-  // Bind this early (instead of inside `useEffect`), because the `SNIPPET_RENDERED` event
-  // is triggered *during* the rendering process, not after. We have to use the ref
-  // to ensure we don't end up calling setState outside the effect though.
-  channel.on(SNIPPET_RENDERED, handleSnippetRendered);
-
   useEffect(() => {
-    const current = sourcesRef.current || {};
-    if (!deepEqual(sources, current)) {
-      setSources(current);
-    }
+    const handleSnippetRendered = (id: StoryId, newItem: SourceItem) => {
+      if (newItem !== sources[id]) {
+        const newSources = { ...sources, [id]: newItem };
+
+        if (!deepEqual(sources, newSources)) {
+          setSources(newSources);
+        }
+      }
+    };
+
+    channel.on(SNIPPET_RENDERED, handleSnippetRendered);
 
     return () => channel.off(SNIPPET_RENDERED, handleSnippetRendered);
   });

addons/docs/src/frameworks/angular/sourceDecorator.ts

@@ -1,4 +1,4 @@
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { PartialStoryFn } from '@storybook/csf';
 import { StoryContext, AngularFramework } from '@storybook/angular';
 import { computesTemplateSourceFromComponent } from '@storybook/angular/renderer';
@@ -44,20 +44,21 @@ export const sourceDecorator = (
 
   const { component, argTypes } = context;
 
+  let toEmit: string;
+  useEffect(() => {
+    if (toEmit) channel.emit(SNIPPET_RENDERED, context.id, prettyUp(template));
+  });
+
   if (component && !userDefinedTemplate) {
     const source = computesTemplateSourceFromComponent(component, props, argTypes);
 
     // We might have a story with a Directive or Service defined as the component
     // In these cases there might exist a template, even if we aren't able to create source from component
     if (source || template) {
-      channel.emit(SNIPPET_RENDERED, context.id, prettyUp(source || template));
+      toEmit = prettyUp(source || template);
     }
-    return story;
-  }
-
-  if (template) {
-    channel.emit(SNIPPET_RENDERED, context.id, prettyUp(template));
-    return story;
+  } else if (template) {
+    toEmit = prettyUp(template);
   }
 
   return story;

addons/docs/src/frameworks/html/sourceDecorator.ts

@@ -1,5 +1,5 @@
 /* global window */
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { ArgsStoryFn, PartialStoryFn, StoryContext } from '@storybook/csf';
 import dedent from 'ts-dedent';
 import { HtmlFramework } from '@storybook/html';
@@ -40,11 +40,13 @@ export function sourceDecorator(
     ? (context.originalStoryFn as ArgsStoryFn<HtmlFramework>)(context.args, context)
     : storyFn();
 
+  let source: string;
   if (typeof story === 'string' && !skipSourceRender(context)) {
-    const source = applyTransformSource(story, context);
-
-    addons.getChannel().emit(SNIPPET_RENDERED, context.id, source);
+    source = applyTransformSource(story, context);
   }
+  useEffect(() => {
+    if (source) addons.getChannel().emit(SNIPPET_RENDERED, context.id, source);
+  });
 
   return story;
 }

addons/docs/src/frameworks/react/jsxDecorator.tsx

@@ -3,7 +3,7 @@ import reactElementToJSXString, { Options } from 'react-element-to-jsx-string';
 import dedent from 'ts-dedent';
 import deprecate from 'util-deprecate';
 
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { StoryContext, ArgsStoryFn, PartialStoryFn } from '@storybook/csf';
 import { logger } from '@storybook/client-logger';
 import { ReactFramework } from '@storybook/react';
@@ -175,11 +175,17 @@ export const jsxDecorator = (
   storyFn: PartialStoryFn<ReactFramework>,
   context: StoryContext<ReactFramework>
 ) => {
+  const skip = skipJsxRender(context);
   const story = storyFn();
 
+  let jsx = '';
+  useEffect(() => {
+    if (!skip) channel.emit(SNIPPET_RENDERED, (context || {}).id, jsx);
+  });
+
   // We only need to render JSX if the source block is actually going to
   // consume it. Otherwise it's just slowing us down.
-  if (skipJsxRender(context)) {
+  if (skip) {
     return story;
   }
 
@@ -197,13 +203,10 @@ export const jsxDecorator = (
 
   const sourceJsx = mdxToJsx(storyJsx);
 
-  let jsx = '';
   const rendered = renderJsx(sourceJsx, options);
   if (rendered) {
     jsx = applyTransformSource(rendered, options, context);
   }
 
-  channel.emit(SNIPPET_RENDERED, (context || {}).id, jsx);
-
   return story;
 };

addons/docs/src/frameworks/svelte/sourceDecorator.ts

@@ -1,4 +1,4 @@
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { ArgTypes, Args, StoryContext, AnyFramework } from '@storybook/csf';
 
 import { SourceType, SNIPPET_RENDERED } from '../../shared';
@@ -145,9 +145,17 @@ function getWrapperProperties(component: any) {
  * @param context  StoryContext
  */
 export const sourceDecorator = (storyFn: any, context: StoryContext<AnyFramework>) => {
+  const skip = skipSourceRender(context);
   const story = storyFn();
 
-  if (skipSourceRender(context)) {
+  let source: string;
+  useEffect(() => {
+    if (!skip && source) {
+      channel.emit(SNIPPET_RENDERED, (context || {}).id, source);
+    }
+  });
+
+  if (skip) {
     return story;
   }
 
@@ -161,11 +169,7 @@ export const sourceDecorator = (storyFn: any, context: StoryContext<AnyFramework
     component = parameters.component;
   }
 
-  const source = generateSvelteSource(component, args, context?.argTypes, slotProperty);
-
-  if (source) {
-    channel.emit(SNIPPET_RENDERED, (context || {}).id, source);
-  }
+  source = generateSvelteSource(component, args, context?.argTypes, slotProperty);
 
   return story;
 };

addons/docs/src/frameworks/web-components/sourceDecorator.ts

@@ -1,7 +1,7 @@
 /* global window */
 import { render } from 'lit-html';
 import { ArgsStoryFn, PartialStoryFn, StoryContext } from '@storybook/csf';
-import { addons } from '@storybook/addons';
+import { addons, useEffect } from '@storybook/addons';
 import { WebComponentsFramework } from '@storybook/web-components';
 
 import { SNIPPET_RENDERED, SourceType } from '../../shared';
@@ -37,11 +37,14 @@ export function sourceDecorator(
     ? (context.originalStoryFn as ArgsStoryFn<WebComponentsFramework>)(context.args, context)
     : storyFn();
 
+  let source: string;
+  useEffect(() => {
+    if (source) addons.getChannel().emit(SNIPPET_RENDERED, context.id, source);
+  });
   if (!skipSourceRender(context)) {
     const container = window.document.createElement('div');
     render(story, container);
-    const source = applyTransformSource(container.innerHTML.replace(/<!---->/g, ''), context);
-    if (source) addons.getChannel().emit(SNIPPET_RENDERED, context.id, source);
+    source = applyTransformSource(container.innerHTML.replace(/<!---->/g, ''), context);
   }
 
   return story;