From 31483d244248353260b0a9cdf71d496b6ac7a05b Mon Sep 17 00:00:00 2001
From: Yossi Saadi <yossisa@monday.com>
Date: Tue, 16 Apr 2024 11:56:58 +0300
Subject: [PATCH] chore(withLiveEdit): should not throw errors on stories that
 do not include LiveEdit (#2070)

---
 .../decorators/withLiveEdit/withLiveEdit.tsx  | 56 ++++++++++---------
 1 file changed, 29 insertions(+), 27 deletions(-)

diff --git a/packages/core/src/storybook/decorators/withLiveEdit/withLiveEdit.tsx b/packages/core/src/storybook/decorators/withLiveEdit/withLiveEdit.tsx
index 3315b93210..c784c33273 100644
--- a/packages/core/src/storybook/decorators/withLiveEdit/withLiveEdit.tsx
+++ b/packages/core/src/storybook/decorators/withLiveEdit/withLiveEdit.tsx
@@ -1,5 +1,5 @@
 import { Decorator, Parameters as StorybookParameters } from "@storybook/react";
-import { useRef, useState } from "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";
@@ -21,25 +21,24 @@ type Parameters = StorybookParameters & {
   };
 };
 
-function getInitialCodeValue(code: string): string {
+function getInitialCodeValue(code: string, shouldPrintError: boolean): string {
   try {
     return formatCode(code);
   } catch (e) {
-    console.error(e);
+    if (shouldPrintError) {
+      console.error(e);
+    }
     return code;
   }
 }
 
 const withLiveEdit: Decorator = (Story, { id, parameters, viewMode }: Parameters) => {
   const scope = { ...globalScope, ...parameters.docs?.liveEdit?.scope };
-  const shouldAllowLiveEdit = viewMode === "docs" && parameters.docs?.liveEdit?.isEnabled;
-
-  if (!shouldAllowLiveEdit) {
-    return <Story />;
-  }
+  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));
+  const [code, setCode] = useState<string>(getInitialCodeValue(originalCode.current, shouldAllowLiveEdit));
   const [dirty, setDirty] = useState<boolean>(false);
 
   const handleChange = (newVal: string) => {
@@ -47,6 +46,10 @@ const withLiveEdit: Decorator = (Story, { id, parameters, viewMode }: Parameters
     setDirty(true);
   };
 
+  if (!shouldAllowLiveEdit) {
+    return <Story />;
+  }
+
   return (
     <>
       {dirty ? (
@@ -59,24 +62,23 @@ const withLiveEdit: Decorator = (Story, { id, parameters, viewMode }: Parameters
       ) : (
         <Story />
       )}
-      {document.getElementById(id) &&
-        createPortal(
-          <LiveEditor
-            placeholder={"Insert code here"}
-            code={code}
-            theme={vscodeDark}
-            extensions={[langs.tsx()]}
-            style={{ fontSize: 13 }}
-            onChange={handleChange}
-            setup={{
-              lineNumbers: false,
-              foldGutter: false,
-              highlightActiveLine: false,
-              autocompletion: false
-            }}
-          />,
-          document.getElementById(id)
-        )}
+      {createPortal(
+        <LiveEditor
+          placeholder={"Insert code here"}
+          code={code}
+          theme={vscodeDark}
+          extensions={[langs.tsx()]}
+          style={{ fontSize: 13 }}
+          onChange={handleChange}
+          setup={{
+            lineNumbers: false,
+            foldGutter: false,
+            highlightActiveLine: false,
+            autocompletion: false
+          }}
+        />,
+        canvasEditorContainer
+      )}
     </>
   );
 };