fix: improve heading-analyzer reliability

Use new status value from usePuck with a MutationObserver.
measuredco · Dec 27, 2024 · 21b118e · 21b118e
1 parent ca5f369
commit 21b118e
Showing 1 changed file with 23 additions and 30 deletions.
diff --git a/packages/plugin-heading-analyzer/src/HeadingAnalyzer.tsx b/packages/plugin-heading-analyzer/src/HeadingAnalyzer.tsx
@@ -22,25 +22,20 @@ const ReactFromJSON =
   (ReactFromJSONModule as unknown as { default: typeof ReactFromJSONModule })
     .default || ReactFromJSONModule;
 
-const dataAttr = "data-puck-heading-analyzer-id";
-
-const getOutline = ({
-  addDataAttr = false,
-  frame,
-}: { addDataAttr?: boolean; frame?: Element | Document } = {}) => {
+const getOutline = ({ frame }: { frame?: Element | Document } = {}) => {
   const headings = frame?.querySelectorAll("h1,h2,h3,h4,h5,h6") || [];
 
-  const _outline: { rank: number; text: string; analyzeId: string }[] = [];
+  const _outline: {
+    rank: number;
+    text: string;
+    element: HTMLElement;
+  }[] = [];
 
   headings.forEach((item, i) => {
-    if (addDataAttr) {
-      item.setAttribute(dataAttr, i.toString());
-    }
-
     _outline.push({
       rank: parseInt(item.tagName.split("H")[1]),
       text: item.textContent!,
-      analyzeId: i.toString(),
+      element: item as HTMLElement,
     });
   });
 
@@ -53,10 +48,11 @@ type Block = {
   children?: Block[];
   missing?: boolean;
   analyzeId?: string;
+  element?: HTMLElement;
 };
 
 function buildHierarchy(frame: Element | Document): Block[] {
-  const headings = getOutline({ addDataAttr: true, frame });
+  const headings = getOutline({ frame });
 
   const root = { rank: 0, children: [], text: "" }; // Placeholder root node
   let path: Block[] = [root];
@@ -65,8 +61,8 @@ function buildHierarchy(frame: Element | Document): Block[] {
     const node: Block = {
       rank: heading.rank,
       text: heading.text,
-      analyzeId: heading.analyzeId,
       children: [],
+      element: heading.element,
     };
 
     // When encountering an h1, reset the path to only the root
@@ -102,24 +98,25 @@ function buildHierarchy(frame: Element | Document): Block[] {
 export const HeadingAnalyzer = () => {
   const { appState } = usePuck();
   const [hierarchy, setHierarchy] = useState<Block[]>([]);
-  const [firstRender, setFirstRender] = useState(true);
 
   // Re-render when content changes
   useEffect(() => {
     const frame = getFrame();
 
     if (!frame) return;
 
-    // We need to delay to allow remainder of page to render first
-    if (firstRender) {
-      setTimeout(() => {
-        setHierarchy(buildHierarchy(frame));
-        setFirstRender(false);
-      }, 100);
-    } else {
+    setHierarchy(buildHierarchy(frame));
+
+    const observer = new MutationObserver(() => {
       setHierarchy(buildHierarchy(frame));
-    }
-  }, [appState.data]);
+    });
+
+    observer.observe(frame.body, { subtree: true, childList: true });
+
+    return () => {
+      observer.disconnect();
+    };
+  }, [appState.ui.status]);
 
   return (
     <div className={getClassName()}>
@@ -153,16 +150,12 @@ export const HeadingAnalyzer = () => {
                   <small
                     className={getClassNameItem({ missing: props.missing })}
                     onClick={
-                      typeof props.analyzeId == "undefined"
+                      typeof props.element == "undefined"
                         ? undefined
                         : (e) => {
                             e.stopPropagation();
 
-                            const frame = getFrame();
-
-                            const el = frame?.querySelector(
-                              `[${dataAttr}="${props.analyzeId}"]`
-                            ) as HTMLElement;
+                            const el = props.element;
 
                             const oldStyle = { ...el.style };