feat(withLiveEdit): apply decorators from within self CSF module of a…

… story (#2077)

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

@@ -0,0 +1,21 @@
+import React from "react";
+import { LiveProvider } from "react-live";
+import LivePreview from "../../../components/live-preview/LivePreview";
+import useApplyDecorators from "../hooks/useApplyDecorators";
+import { LiveContentProps } from "./LiveContent.types";
+import styles from "./LiveContent.module.scss";
+
+const LiveContent = ({ code, scope, decorators, context }: LiveContentProps) => {
+  const content: React.JSX.Element = (
+    <>
+      <div className={styles.modifiedVersionIndicator}>Modified Version</div>
+      <LiveProvider code={code} scope={scope} enableTypeScript>
+        <LivePreview />
+      </LiveProvider>
+    </>
+  );
+
+  return <>{useApplyDecorators(decorators || [], content, context)}</>;
+};
+
+export default LiveContent;

packages/core/src/storybook/decorators/withLiveEdit/LiveContent/LiveContent.types.ts

@@ -0,0 +1,8 @@
+import { Decorator, StoryContext } from "@storybook/react";
+
+export interface LiveContentProps {
+  code: string;
+  scope: Record<string, unknown>;
+  decorators: Decorator[];
+  context: StoryContext;
+}

packages/core/src/storybook/decorators/withLiveEdit/hooks/useApplyDecorators.ts

@@ -0,0 +1,18 @@
+import React, { useMemo } from "react";
+import { StoryContext, Decorator } from "@storybook/react";
+
+const useApplyDecorators = (decorators: Decorator[], component: React.ReactElement, context: StoryContext) => {
+  return useMemo(() => {
+    let decoratedComponent = () => component;
+
+    // recursively apply decorators to the component
+    decorators.forEach(decorator => {
+      const currentComponent = decoratedComponent;
+      decoratedComponent = () => decorator(currentComponent, context);
+    });
+
+    return decoratedComponent();
+  }, [decorators, component, context]);
+};
+
+export default useApplyDecorators;

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

@@ -1,26 +1,17 @@
-import { Decorator, Parameters as StorybookParameters } from "@storybook/react";
+import { Decorator, StoryContext } from "@storybook/react";
 import { useMemo, useRef, useState } from "react";
 import { vscodeDark } from "@uiw/codemirror-theme-vscode";
 import { langs } from "@uiw/codemirror-extensions-langs";
 import * as VibeComponents from "../../../components";
 import * as VibeComponentsNext from "../../../next";
 import * as VibeIcons from "../../../components/Icon/Icons";
 import { createPortal } from "react-dom";
-import { LiveProvider } from "react-live";
-import LivePreview from "../../components/live-preview/LivePreview";
-import styles from "./withLiveEdit.module.scss";
 import LiveEditor from "../../components/live-editor/LiveEditor";
 import { formatCode } from "./prettier-utils";
+import LiveContent from "./LiveContent/LiveContent";
 
 const globalScope = { ...VibeComponents, VibeIcons, VibeNext: VibeComponentsNext };
 
-type Parameters = StorybookParameters & {
-  liveEdit?: {
-    scope?: Record<string, unknown>;
-    isEnabled?: boolean;
-  };
-};
-
 function getInitialCodeValue(code: string, shouldPrintError: boolean): string {
   try {
     return formatCode(code);
@@ -32,7 +23,8 @@ function getInitialCodeValue(code: string, shouldPrintError: boolean): string {
   }
 }
 
-const withLiveEdit: Decorator = (Story, { id, parameters, viewMode }: Parameters) => {
+const withLiveEdit: Decorator = (Story, context: StoryContext) => {
+  const { id, parameters, viewMode, moduleExport } = context;
   const scope = { ...globalScope, ...parameters.docs?.liveEdit?.scope };
   const canvasEditorContainer = useMemo(() => document.getElementById(id), [id]);
   const shouldAllowLiveEdit = viewMode === "docs" && parameters.docs?.liveEdit?.isEnabled && !!canvasEditorContainer;
@@ -53,12 +45,7 @@ const withLiveEdit: Decorator = (Story, { id, parameters, viewMode }: Parameters
   return (
     <>
       {dirty ? (
-        <>
-          <div className={styles.modifiedVersionIndicator}>Modified Version</div>
-          <LiveProvider code={code} scope={scope} enableTypeScript>
-            <LivePreview />
-          </LiveProvider>
-        </>
+        <LiveContent code={code} scope={scope} decorators={moduleExport.decorators || []} context={context} />
       ) : (
         <Story />
       )}