feat: ability to show source for Playground

.vscode/launch.json

@@ -139,7 +139,7 @@
       "program": "${workspaceFolder}/node_modules/.bin/jest",
       "cwd": "${workspaceFolder}/core/instrument",
       "args": [
-        "mdx-async",
+        "mdx-stories",
       ],
       "console": "integratedTerminal",
       "internalConsoleOptions": "neverOpen",

core/instrument/src/babel/extract-attributes.ts

@@ -92,3 +92,25 @@ export const extractAttributes = (
   }
   return undefined;
 };
+
+export const collectAttributes = (node: any): Record<string, string> => {
+  return node.attributes.reduce(
+    (acc: Record<string, unknown>, attribute: any) => {
+      if (!attribute.value) {
+        //console.log(attribute);
+      } else if (attribute.value.type === 'StringLiteral') {
+        return {
+          ...acc,
+          [attribute.name.name]: attribute.value.value,
+        };
+      } else if (attribute.value.type === 'JSXExpressionContainer') {
+        return {
+          ...acc,
+          [attribute.name.name]: extractAttributes(attribute.value.expression),
+        };
+      }
+      return acc;
+    },
+    {},
+  );
+};

core/instrument/src/babel/extract-function-parameters.ts

@@ -131,49 +131,3 @@ export const extractFunctionParameters = (
     path.skip();
   },
 });
-
-const adjustCodeLocation = (
-  loc: CodeLocation,
-  line: number,
-  chars: number,
-): CodeLocation => {
-  return {
-    start: {
-      line: loc.start.line,
-      column:
-        line === loc.start.line
-          ? Math.max(0, loc.start.column - chars)
-          : loc.start.column,
-    },
-    end: {
-      line: loc.end.line,
-      column:
-        line === loc.end.line
-          ? Math.max(0, loc.end.column - chars)
-          : loc.end.column,
-    },
-  };
-};
-
-// adjust source location, when timmiong white spaces
-export const adjustArgumentsLocation = (
-  args: StoryArguments,
-  line: number,
-  chars: number,
-): StoryArguments =>
-  args.map(arg => {
-    const ret = {
-      ...arg,
-      loc: arg.loc ? adjustCodeLocation(arg.loc, line, chars) : undefined,
-      value: Array.isArray(arg.value)
-        ? adjustArgumentsLocation(arg.value, line, chars)
-        : arg.value,
-    };
-    if (arg.usage) {
-      ret.usage = arg.usage.map(usage => ({
-        ...usage,
-        loc: adjustCodeLocation(usage.loc, line, chars),
-      }));
-    }
-    return ret;
-  });

core/instrument/src/babel/mdx-stories.ts

@@ -7,15 +7,12 @@ import {
 import camelCase from 'camelcase';
 import { File } from '@babel/types';
 import traverse from '@babel/traverse';
-import {
-  extractFunctionParameters,
-  adjustArgumentsLocation,
-} from './extract-function-parameters';
+import { extractFunctionParameters } from './extract-function-parameters';
 import { followStoryImport } from './follow-imports';
-import { extractAttributes } from './extract-attributes';
+import { collectAttributes } from './extract-attributes';
 import { sourceLocation } from '../misc/source-location';
 import { ParseStorieReturnType, InstrumentOptions } from '../types';
-
+import { extractSource } from '../misc/source-extract';
 import { componentsFromParams } from '../misc/component-attributes';
 
 type PartialStore = Required<
@@ -40,30 +37,6 @@ export const extractMDXStories: (
   _options: Required<InstrumentOptions>,
   { source, filePath }: SourceFile,
 ): PartialStore | undefined => {
-  const collectAttributes = (node: any): Record<string, string> => {
-    return node.attributes.reduce(
-      (acc: Record<string, unknown>, attribute: any) => {
-        if (!attribute.value) {
-          //console.log(attribute);
-        } else if (attribute.value.type === 'StringLiteral') {
-          return {
-            ...acc,
-            [attribute.name.name]: attribute.value.value,
-          };
-        } else if (attribute.value.type === 'JSXExpressionContainer') {
-          return {
-            ...acc,
-            [attribute.name.name]: extractAttributes(
-              attribute.value.expression,
-            ),
-          };
-        }
-        return acc;
-      },
-      {},
-    );
-  };
-
   let components: { [key: string]: string | undefined } = {};
   const collectComponent = (attributes: any) => {
     const attrComponents = componentsFromParams(attributes);
@@ -186,38 +159,8 @@ export const extractMDXStories: (
                 );
                 // adjust for source code wraping issues
                 const storySource = getASTSource(source, loc);
-                const storyLines = storySource?.split('\n');
-                //remove as many spaces as there are in the line with the smallest amount of white spaces (but still some)
-                const whiteSpaces = storyLines?.reduce((spaces, line) => {
-                  if (line.substring(0, 1) === ' ' && line.trim() !== '') {
-                    const startSpaces = line.search(/\S/);
-                    if (startSpaces && (spaces === 0 || startSpaces < spaces)) {
-                      return startSpaces;
-                    }
-                  }
-                  return spaces;
-                }, 0);
                 story.loc = sourceLocation(loc);
-                story.source =
-                  storyLines && whiteSpaces
-                    ? storyLines
-                        .map((line, lineIndex) => {
-                          if (line.search(/\S/) >= whiteSpaces) {
-                            //adjust location of arguments usage
-                            if (story.arguments) {
-                              story.arguments = adjustArgumentsLocation(
-                                story.arguments,
-                                lineIndex,
-                                whiteSpaces,
-                              );
-                            }
-                            return line.substr(whiteSpaces);
-                          }
-                          return line;
-                        })
-                        .join('\n')
-                        .trim()
-                    : storySource;
+                story.source = extractSource(storySource, story);
               }
 
               store.stories[id] = story;
@@ -254,7 +197,6 @@ export const extractMDXStories: (
             collectComponent(attributes);
             break;
         }
-        // console.log(node.name.name, attributes);
       }
     },
   });

core/instrument/src/babel/transform-ast-tree.ts

@@ -1,9 +1,62 @@
 import { TraverseOptions } from '@babel/traverse';
-
 const babel = require('@babel/core');
+import { CodeLocation, getASTSource } from '@component-controls/core';
+import { extractSource } from '../misc/source-extract';
+import { collectAttributes } from './extract-attributes';
 
-export const transformTree = (): TraverseOptions => {
+export const transformMDXAttributes = (source: string): TraverseOptions => {
   return {
+    JSXElement: (path: any) => {
+      const node = path.node.openingElement;
+      switch (node.name.name) {
+        case 'Playground':
+          const playground = collectAttributes(node);
+          // extract playground source of not already a prop
+          if (!playground.source) {
+            const hasStory = path.node.children.some(
+              (el: any) =>
+                el.type === 'JSXElement' &&
+                el.openingElement?.name?.name === 'Story',
+            );
+            if (!hasStory) {
+              const loc: CodeLocation = path.node.children.length
+                ? {
+                    start: path.node.children[0].loc.start,
+                    end:
+                      path.node.children[path.node.children.length - 1].loc.end,
+                  }
+                : path.node.loc;
+              const rawSource = getASTSource(source, loc);
+              const prettiedSource = extractSource(rawSource);
+              node.attributes.push({
+                type: 'JSXAttribute',
+                name: {
+                  type: 'JSXIdentifier',
+                  name: 'source',
+                },
+                value: {
+                  type: 'JSXExpressionContainer',
+                  expression: {
+                    type: 'TemplateLiteral',
+                    expressions: [],
+                    quasis: [
+                      {
+                        type: 'TemplateElement',
+                        value: {
+                          raw: prettiedSource,
+                          cooked: prettiedSource,
+                        },
+                        tail: true,
+                      },
+                    ],
+                  },
+                },
+              });
+            }
+          }
+          break;
+      }
+    },
     JSXAttribute: (path: any) => {
       const node = path.node;
       switch (node.name?.name) {

core/instrument/src/index.ts

@@ -14,7 +14,7 @@ import {
 
 import { extractCSFStories } from './babel/esm-stories';
 import { extractMDXStories } from './babel/mdx-stories';
-import { transformTree } from './babel/transform-ast-tree';
+import { transformMDXAttributes } from './babel/transform-ast-tree';
 import { extractStoreComponent } from './babel/extract-component';
 import { packageInfo } from './misc/package-info';
 import { extractStoryExports } from './misc/mdx-exports';
@@ -61,7 +61,7 @@ const parseSource = async (
   filePath: string,
   options: Required<InstrumentOptions>,
 ): Promise<ParseStorieReturnType | undefined> => {
-  const source = await prettify(code, filePath, options.prettier);
+  const source = await prettify(code, options.prettier, filePath);
   const ast = parser.parse(source, options.parser);
 
   const store = traverseFn(ast, options, { source, filePath });
@@ -158,16 +158,20 @@ export const parseStories = async (
   } = mergedOptions.mdx;
   if (test && filePath.match(test)) {
     const { data, content } = matter(source);
-    const mdxParsed = await mdx(content, {
+    code = await mdx(content, {
       filepath: filePath,
       ...otherMDXOptions,
     });
-    const ast = parser.parse(mdxParsed, mergedOptions.parser) as any;
-    traverse(ast, transformTree());
-    ({ code } = generate(ast, {
-      retainFunctionParens: true,
-      retainLines: true,
-    }));
+    const ast = parser.parse(code, mergedOptions.parser) as any;
+
+    if (transformMDX) {
+      //second pass transform - inject any necessary attributes
+      traverse(ast, transformMDXAttributes(code));
+      ({ code } = generate(ast, {
+        retainFunctionParens: true,
+        retainLines: false,
+      }));
+    }
     const store = await parseSource(
       code,
       extractMDXStories(data),

core/instrument/src/misc/prettify.ts

@@ -4,8 +4,8 @@ import { PrettierOptions } from '../types';
 
 export const prettify = async (
   code: string,
-  filePath: string,
   options: PrettierOptions | false,
+  filePath?: string,
 ): Promise<string> => {
   if (options !== false) {
     const { resolveConfigOptions, ...otherOptions } = options || {};

core/instrument/src/misc/source-extract.ts

@@ -0,0 +1,91 @@
+import { CodeLocation, StoryArguments, Story } from '@component-controls/core';
+
+const adjustCodeLocation = (
+  loc: CodeLocation,
+  line: number,
+  chars: number,
+): CodeLocation => {
+  return {
+    start: {
+      line: loc.start.line,
+      column:
+        line === loc.start.line
+          ? Math.max(0, loc.start.column - chars)
+          : loc.start.column,
+    },
+    end: {
+      line: loc.end.line,
+      column:
+        line === loc.end.line
+          ? Math.max(0, loc.end.column - chars)
+          : loc.end.column,
+    },
+  };
+};
+
+// adjust source location, when timmiong white spaces
+export const adjustArgumentsLocation = (
+  args: StoryArguments,
+  line: number,
+  chars: number,
+): StoryArguments =>
+  args.map(arg => {
+    const ret = {
+      ...arg,
+      loc: arg.loc ? adjustCodeLocation(arg.loc, line, chars) : undefined,
+      value: Array.isArray(arg.value)
+        ? adjustArgumentsLocation(arg.value, line, chars)
+        : arg.value,
+    };
+    if (arg.usage) {
+      ret.usage = arg.usage.map(usage => ({
+        ...usage,
+        loc: adjustCodeLocation(usage.loc, line, chars),
+      }));
+    }
+    return ret;
+  });
+
+export const removeMDXAttributes = (source: string): string =>
+  source.replace(/mdxType="[^"]*"\s?/, '');
+export const extractSource = (
+  source?: string,
+  story?: Story,
+): string | undefined => {
+  if (!source) {
+    return undefined;
+  }
+  const code = removeMDXAttributes(source);
+  const lines = code.split('\n');
+  //remove as many spaces as there are in the line with the smallest amount of white spaces (but still some)
+  const whiteSpacesList = lines
+    .map(line => {
+      if (line.substring(0, 1) === ' ' && line.trim() !== '') {
+        return line.search(/\S/);
+      }
+      return 999999;
+    })
+    .sort((a: number, b: number) => a - b);
+  const whiteSpaces = whiteSpacesList.length > 1 && whiteSpacesList[1];
+  return whiteSpaces
+    ? lines
+        .map((line, lineIndex) => {
+          const spaces = line.search(/\S/);
+          const adjustSpaces = Math.min(spaces, whiteSpaces);
+          if (adjustSpaces > 0) {
+            //adjust location of arguments usage
+            if (story?.arguments) {
+              story.arguments = adjustArgumentsLocation(
+                story.arguments,
+                lineIndex,
+                adjustSpaces,
+              );
+            }
+            return line.substr(adjustSpaces);
+          }
+          return line;
+        })
+        .join('\n')
+        .trim()
+    : code;
+};