Addon-docs: Automatic source selection based on story type

addons/docs/src/blocks/Source.tsx

@@ -4,12 +4,38 @@ import {
   SourceError,
   SourceProps as PureSourceProps,
 } from '@storybook/components';
+import { StoryId } from '@storybook/api';
+import { logger } from '@storybook/client-logger';
+
 import { DocsContext, DocsContextProps } from './DocsContext';
 import { SourceContext, SourceContextProps } from './SourceContainer';
 import { CURRENT_SELECTION } from './types';
 
 import { enhanceSource } from './enhanceSource';
 
+enum SourceType {
+  /**
+   * AUTO is the default
+   *
+   * Use the CODE logic if:
+   * - the user has set a custom source snippet in `docs.source.code` story parameter
+   * - the story is not an args-based story
+   *
+   * Use the DYNAMIC rendered snippet if the story is an args story
+   */
+  AUTO = 'auto',
+
+  /**
+   * Render the code extracted by source-loader
+   */
+  CODE = 'code',
+
+  /**
+   * Render dynamically-rendered source snippet from the story's virtual DOM (currently React only)
+   */
+  DYNAMIC = 'dynamic',
+}
+
 interface CommonProps {
   language?: string;
   dark?: boolean;
@@ -32,13 +58,53 @@ type NoneProps = CommonProps;
 
 type SourceProps = SingleSourceProps | MultiSourceProps | CodeProps | NoneProps;
 
+const getSnippet = (
+  storyId: StoryId,
+  sourceContext: SourceContextProps,
+  docsContext: DocsContextProps
+): string => {
+  const { sources } = sourceContext;
+  const { storyStore } = docsContext;
+
+  const snippet = sources && sources[storyId];
+  const data = storyStore?.fromId(storyId);
+
+  if (data) {
+    // Fallback if we can't get the story data for this story
+    logger.warn(`Unable to find source for story ID '${storyId}'`);
+    return snippet || '';
+  }
+
+  const { parameters } = data;
+  // eslint-disable-next-line no-underscore-dangle
+  const isArgsStory = parameters.__isArgsStory;
+  const type = parameters.docs?.source?.type || SourceType.AUTO;
+
+  // if user has hard-coded the snippet, that takes precedence
+  const userCode = parameters.docs?.source?.code;
+  if (userCode) return userCode;
+
+  // if user has explicitly set this as dynamic, use snippet
+  if (type === SourceType.DYNAMIC) {
+    return snippet || '';
+  }
+
+  // if this is an args story and there's a snippet
+  if (type === SourceType.AUTO && snippet && isArgsStory) {
+    return snippet;
+  }
+
+  // otherwise, use the source code logic
+  const enhanced = enhanceSource(data) || data.parameters;
+  return enhanced?.docs?.source?.code || '';
+};
+
 export const getSourceProps = (
   props: SourceProps,
   docsContext: DocsContextProps,
   sourceContext: SourceContextProps
 ): PureSourceProps => {
-  const { id: currentId, storyStore } = docsContext;
-  const { sources } = sourceContext;
+  const { id: currentId } = docsContext;
 
   const codeProps = props as CodeProps;
   const singleProps = props as SingleSourceProps;
@@ -50,16 +116,7 @@ export const getSourceProps = (
       singleProps.id === CURRENT_SELECTION || !singleProps.id ? currentId : singleProps.id;
     const targetIds = multiProps.ids || [targetId];
     source = targetIds
-      .map((sourceId) => {
-        const snippet = sources && sources[sourceId];
-        if (snippet) return snippet;
-        if (storyStore) {
-          const data = storyStore.fromId(sourceId);
-          const enhanced = data && (enhanceSource(data) || data.parameters);
-          return enhanced?.docs?.source?.code || '';
-        }
-        return '';
-      })
+      .map((storyId) => getSnippet(storyId, sourceContext, docsContext))
       .join('\n\n');
   }
   return source

examples/official-storybook/stories/addon-docs/source.stories.tsx

@@ -0,0 +1,32 @@
+import React from 'react';
+import Button from '../../components/TsButton';
+
+export default {
+  title: 'Addons/Docs/Source',
+  component: Button,
+  argTypes: {
+    children: { control: 'text' },
+    type: { control: 'text' },
+  },
+  parameters: {
+    chromatic: { disable: true },
+  },
+};
+
+const Template = (args) => <Button {...args} />;
+
+export const Basic = Template.bind({});
+Basic.args = {
+  children: 'basic',
+  somethingElse: { a: 2 },
+};
+
+export const NoArgs = () => <Button>no args</Button>;
+
+export const ForceCodeSource = Template.bind({});
+ForceCodeSource.args = { ...Basic.args };
+ForceCodeSource.parameters = { docs: { source: { type: 'code' } } };
+
+export const CustomSource = Template.bind({});
+CustomSource.args = { ...Basic.args };
+CustomSource.parameters = { docs: { source: { code: 'custom source' } } };