feat(withLiveEdit): parse render attribute with ast instead of with regex for variety of cases

packages/core/package.json

@@ -126,6 +126,7 @@
   "devDependencies": {
     "@babel/core": "^7.23.2",
     "@babel/eslint-parser": "^7.16.5",
+    "@babel/parser": "^7.24.4",
     "@babel/plugin-proposal-class-properties": "^7.16.5",
     "@babel/plugin-proposal-private-methods": "^7.18.6",
     "@babel/plugin-proposal-private-property-in-object": "^7.21.11",
@@ -135,6 +136,8 @@
     "@babel/preset-env": "^7.16.5",
     "@babel/preset-react": "^7.16.5",
     "@babel/preset-typescript": "^7.23.3",
+    "@babel/standalone": "^7.24.4",
+    "@babel/types": "^7.24.0",
     "@hot-loader/react-dom": "^16.13.0",
     "@jest/transform": "^26.6.2",
     "@mdx-js/loader": "^2.0.0-rc.2",
@@ -170,6 +173,7 @@
     "@testing-library/react-hooks": "^7.0.2",
     "@testing-library/user-event": "^13.5.0",
     "@types/autosize": "^4.0.1",
+    "@types/babel__standalone": "^7.1.7",
     "@types/body-scroll-lock": "^3.1.0",
     "@types/lodash": "^4.14.184",
     "@types/lodash-es": "^4.17.6",

packages/core/src/storybook/decorators/withLiveEdit/utils/parse-csf-utils.ts

@@ -0,0 +1,90 @@
+import { transformFromAst } from "@babel/standalone";
+import { parse } from "@babel/parser";
+import {
+  File,
+  Program,
+  Expression,
+  Statement,
+  ObjectProperty,
+  FunctionExpression,
+  ArrowFunctionExpression,
+  isObjectExpression,
+  isBlockStatement,
+  isExpressionStatement,
+  isObjectProperty,
+  isIdentifier,
+  isArrowFunctionExpression,
+  isFunctionExpression
+} from "@babel/types";
+
+const parseCsfToAst = (csfString: string) => {
+  return parse(csfString, {
+    plugins: ["typescript", "jsx"],
+    sourceType: "module"
+  });
+};
+
+function findRenderProperty(body: Statement[]): ObjectProperty {
+  return body
+    .flatMap(node =>
+      isExpressionStatement(node) && isObjectExpression(node.expression) ? node.expression.properties : []
+    )
+    .find(prop => isObjectProperty(prop) && isIdentifier(prop.key) && prop.key.name === "render") as ObjectProperty;
+}
+
+function prepareTransformTarget(renderProperty: ObjectProperty): Program {
+  const { value } = renderProperty;
+  const isFunction = isFunctionExpression(value) || isArrowFunctionExpression(value);
+  if (!isFunction) {
+    // e.g. render: something
+    const expression = value as Expression;
+    return {
+      type: "Program",
+      body: [{ type: "ExpressionStatement", expression }],
+      directives: [],
+      sourceType: "module"
+    };
+  }
+
+  // e.g. render: () => {...}, render: () => (...), render()
+  const { body } = value as FunctionExpression | ArrowFunctionExpression;
+
+  return {
+    type: "Program",
+    body: [
+      {
+        type: "ExpressionStatement",
+        expression: isBlockStatement(body) ? value : body
+      }
+    ],
+    directives: [],
+    sourceType: "module"
+  };
+}
+
+function transformAstToCode(ast: File): string {
+  // docs says transformFromAst is void, but it is not because of a breaking change
+  const { code } = transformFromAst(ast, null, { presets: [] }) as unknown as { code: string };
+  return code;
+}
+
+export function extractRenderAttributeFromCsf(csfString: string) {
+  const originalAst = parseCsfToAst(csfString);
+
+  const renderProperty = findRenderProperty(originalAst.program.body);
+  if (!renderProperty) {
+    return null;
+  }
+
+  const transformTarget = prepareTransformTarget(renderProperty);
+  if (!transformTarget) {
+    return null;
+  }
+
+  const transformedAst: File = {
+    type: "File",
+    program: transformTarget
+  };
+
+  return transformAstToCode(transformedAst);
+}

packages/core/src/storybook/decorators/withLiveEdit/withLiveEdit.tsx

@@ -7,19 +7,22 @@ import * as VibeComponentsNext from "../../../next";
 import * as VibeIcons from "../../../components/Icon/Icons";
 import { createPortal } from "react-dom";
 import LiveEditor from "../../components/live-editor/LiveEditor";
-import { formatCode } from "./prettier-utils";
+import { formatCode } from "./utils/prettier-utils";
 import LiveContent from "./LiveContent/LiveContent";
+import { extractRenderAttributeFromCsf } from "./utils/parse-csf-utils";
 
 const globalScope = { ...VibeComponents, VibeIcons, VibeNext: VibeComponentsNext };
 
-function getInitialCodeValue(code: string, shouldPrintError: boolean): string {
+function getInitialCodeValue(source: string, shouldPrintError: boolean): string {
   try {
+    // need to wrap with parentheses to avoid syntax errors
+    const code = extractRenderAttributeFromCsf(`(${source})`);
     return formatCode(code);
   } catch (e) {
     if (shouldPrintError) {
       console.error(e);
     }
-    return code;
+    return source;
   }
 }
 
@@ -29,8 +32,10 @@ const withLiveEdit: Decorator = (Story, context: StoryContext) => {
   const canvasEditorContainer = useMemo(() => document.getElementById(id), [id]);
   const shouldAllowLiveEdit = viewMode === "docs" && parameters.docs?.liveEdit?.isEnabled && !!canvasEditorContainer;
 
-  const originalCode = useRef<string>(extractCodeFromSource(parameters.docs.source?.originalSource) || "");
-  const [code, setCode] = useState<string>(getInitialCodeValue(originalCode.current, shouldAllowLiveEdit));
+  const originalCode = useRef<string>(
+    getInitialCodeValue(parameters.docs.source?.originalSource || "", shouldAllowLiveEdit)
+  );
+  const [code, setCode] = useState<string>(originalCode.current);
   const [dirty, setDirty] = useState<boolean>(false);
 
   const handleChange = (newVal: string) => {
@@ -70,18 +75,4 @@ const withLiveEdit: Decorator = (Story, context: StoryContext) => {
   );
 };
 
-function extractCodeFromSource(csfSource: string): string {
-  // capture "render:" from the string
-  if (!csfSource) {
-    return "";
-  }
-  const match = csfSource.match(/render:\s*(?:\(\)\s*=>\s*)?([\s\S]*?)(?=\s*,\s*[\w]+:\s*|}$)/);
-
-  if (!match?.[1]) {
-    return "";
-  }
-
-  return match[1].trim();
-}
-
 export default withLiveEdit;

yarn.lock

@@ -284,6 +284,11 @@
   resolved "https://registry.npmjs.org/@babel/parser/-/parser-7.24.1.tgz#1e416d3627393fab1cb5b0f2f1796a100ae9133a"
   integrity sha512-Zo9c7N3xdOIQrNip7Lc9wvRPzlRtovHVE4lkz8WEDr7uYh/GMQhSiIgFxGIArRHYdJE5kxtZjAf8rT0xhdLCzg==
 
+"@babel/parser@^7.24.4":
+  version "7.24.4"
+  resolved "https://registry.npmjs.org/@babel/parser/-/parser-7.24.4.tgz#234487a110d89ad5a3ed4a8a566c36b9453e8c88"
+  integrity sha512-zTvEBcghmeBma9QIGunWevvBAp4/Qu9Bdq+2k0Ot4fVMD6v3dsC9WOcRSKk7tRRyBM/53yKMJko9xOatGQAwSg==
+
 "@babel/plugin-bugfix-safari-id-destructuring-collision-in-function-expression@^7.24.1":
   version "7.24.1"
   resolved "https://registry.npmjs.org/@babel/plugin-bugfix-safari-id-destructuring-collision-in-function-expression/-/plugin-bugfix-safari-id-destructuring-collision-in-function-expression-7.24.1.tgz#b645d9ba8c2bc5b7af50f0fe949f9edbeb07c8cf"
@@ -1102,6 +1107,11 @@
   dependencies:
     regenerator-runtime "^0.14.0"
 
+"@babel/standalone@^7.24.4":
+  version "7.24.4"
+  resolved "https://registry.npmjs.org/@babel/standalone/-/standalone-7.24.4.tgz#9461220fd641a92fff4be19b34fdb9d18e80d37d"
+  integrity sha512-V4uqWeedadiuiCx5P5OHYJZ1PehdMpcBccNCEptKFGPiZIY3FI5f2ClxUl4r5wZ5U+ohcQ+4KW6jX2K6xXzq4Q==
+
 "@babel/template@^7.22.15", "@babel/template@^7.24.0", "@babel/template@^7.3.3":
   version "7.24.0"
   resolved "https://registry.npmjs.org/@babel/template/-/template-7.24.0.tgz#c6a524aa93a4a05d66aaf31654258fae69d87d50"
@@ -4501,7 +4511,7 @@
   resolved "https://registry.npmjs.org/@types/autosize/-/autosize-4.0.3.tgz#77548c35bca4cc6281593228cde9a50112e1f94c"
   integrity sha512-o0ZyU3ePp3+KRbhHsY4ogjc+ZQWgVN5h6j8BHW5RII4cFKi6PEKK9QPAcphJVkD0dGpyFnD3VRR0WMvHVjCv9w==
 
-"@types/babel__core@^7.0.0", "@types/babel__core@^7.1.14", "@types/babel__core@^7.1.7", "@types/babel__core@^7.18.0":
+"@types/babel__core@^7.0.0", "@types/babel__core@^7.1.0", "@types/babel__core@^7.1.14", "@types/babel__core@^7.1.7", "@types/babel__core@^7.18.0":
   version "7.20.5"
   resolved "https://registry.npmjs.org/@types/babel__core/-/babel__core-7.20.5.tgz#3df15f27ba85319caa07ba08d0721889bb39c017"
   integrity sha512-qoQprZvz5wQFJwMDqeseRXWv3rqMvhgpbXFfVyWhbx9X47POIA6i/+dXefEmZKoAgOaTdaIgNSMqMIU61yRyzA==
@@ -4519,6 +4529,13 @@
   dependencies:
     "@babel/types" "^7.0.0"
 
+"@types/babel__standalone@^7.1.7":
+  version "7.1.7"
+  resolved "https://registry.npmjs.org/@types/babel__standalone/-/babel__standalone-7.1.7.tgz#8ab09548a24f54015e7d84a55486148180bc4ace"
+  integrity sha512-4RUJX9nWrP/emaZDzxo/+RYW8zzLJTXWJyp2k78HufG459HCz754hhmSymt3VFOU6/Wy+IZqfPvToHfLuGOr7w==
+  dependencies:
+    "@types/babel__core" "^7.1.0"
+
 "@types/babel__template@*":
   version "7.4.4"
   resolved "https://registry.npmjs.org/@types/babel__template/-/babel__template-7.4.4.tgz#5672513701c1b2199bc6dad636a9d7491586766f"