🪐 Thebe: fix/missing outputs on page load (#125)

* 💡 simplified use of refs * ✨ simplify! all we need are local refs * 🧮 removed render counters * 🪫mb to controls * 🏷 Add changeset
jupyter-book · Jun 14, 2023 · 069b1da · 069b1da
1 parent 6815413
commit 069b1da
Show file tree

Hide file tree

Showing 5 changed files with 93 additions and 115 deletions.
diff --git a/.changeset/twenty-kings-beam.md b/.changeset/twenty-kings-beam.md
@@ -0,0 +1,5 @@
+---
+'@myst-theme/jupyter': patch
+---
+
+Improvements to thebe for initial rendering
diff --git a/packages/jupyter/src/jupyter.tsx b/packages/jupyter/src/jupyter.tsx
@@ -5,27 +5,34 @@ import type { MinifiedOutput } from 'nbtx';
 import { convertToIOutputs } from 'nbtx';
 import { fetchAndEncodeOutputImages } from './convertImages';
 import type { ThebeCore } from 'thebe-core';
-import { useCellRef, useCellRefRegistry, useNotebookCellExecution } from './providers';
+import { useNotebookCellExecution } from './providers';
 import { SourceFileKind } from 'myst-common';
 import { useXRefState } from '@myst-theme/providers';
 import { useThebeLoader } from 'thebe-react';
 
 function ActiveOutputRenderer({ id, data }: { id: string; data: IOutput[] }) {
-  const ref = useCellRef(id);
   const exec = useNotebookCellExecution(id);
+  const ref = useRef<HTMLDivElement>(null);
 
   useEffect(() => {
-    if (!ref?.el || !exec?.cell) return;
+    if (!ref.current || !exec?.cell) {
+      console.debug(`No cell ref available for cell ${exec?.cell?.id}`);
+      return;
+    }
+    if (exec.cell.isAttachedToDOM) {
+      console.debug(`Cell ${exec.cell.id} already attached to DOM`);
+      return;
+    }
     console.debug(`Attaching cell ${exec.cell.id} to DOM at:`, {
-      el: ref.el,
-      connected: ref.el.isConnected,
+      el: ref.current,
+      connected: ref.current.isConnected,
       data,
     });
-    exec.cell.attachToDOM(ref.el);
+    exec.cell.attachToDOM(ref.current);
     exec.cell.render(data);
-  }, [ref?.el, exec?.cell]);
+  }, [ref?.current, exec?.cell]);
 
-  return null;
+  return <div ref={ref} data-thebe-active-ref="true" className="relative" />;
 }
 
 function PassiveOutputRenderer({
@@ -39,67 +46,65 @@ function PassiveOutputRenderer({
   core: ThebeCore;
   kind: SourceFileKind;
 }) {
-  const [cell] = useState(new core.PassiveCellRenderer(id, undefined, undefined));
+  const cell = useRef(new core.PassiveCellRenderer(id, undefined, undefined));
   const ref = useRef<HTMLDivElement>(null);
-  useEffect(() => {
-    cell.render(data, kind === SourceFileKind.Article);
-  }, [data, cell]);
+
   useEffect(() => {
     if (!ref.current) return;
-    cell.attachToDOM(ref.current, true);
+    cell.current.attachToDOM(ref.current, true);
+    cell.current.render(data, kind === SourceFileKind.Article);
   }, [ref]);
+
   return <div ref={ref} data-thebe-passive-ref="true" />;
 }
 
-const MemoPassiveOutputRenderer = React.memo(PassiveOutputRenderer);
+export const JupyterOutputs = React.memo(
+  ({ id, outputs }: { id: string; outputs: MinifiedOutput[] }) => {
+    const { core, load } = useThebeLoader();
+    const { inCrossRef } = useXRefState();
+    const { data, error } = useFetchAnyTruncatedContent(outputs);
+    const [fullOutputs, setFullOutputs] = useState<IOutput[] | null>(null);
+    const exec = useNotebookCellExecution(id);
 
-export const JupyterOutputs = ({ id, outputs }: { id: string; outputs: MinifiedOutput[] }) => {
-  const { core, load } = useThebeLoader();
-  const { inCrossRef } = useXRefState();
-  const { data, error } = useFetchAnyTruncatedContent(outputs);
-  const [loaded, setLoaded] = useState(false);
-  const [fullOutputs, setFullOutputs] = useState<IOutput[] | null>(null);
-  const registry = useCellRefRegistry();
-  const exec = useNotebookCellExecution(id);
+    useEffect(() => {
+      if (core) return;
+      load();
+    }, [core, load]);
 
-  useEffect(() => {
-    if (core) return;
-    load();
-  }, [core, load]);
+    useEffect(() => {
+      if (!data || fullOutputs != null) return;
 
-  useEffect(() => {
-    if (!data || loaded || fullOutputs != null) return;
-    setLoaded(true);
-    fetchAndEncodeOutputImages(data).then((out) => {
-      const compactOutputs = convertToIOutputs(out, {});
-      setFullOutputs(compactOutputs);
-    });
-  }, [id, data, fullOutputs]);
+      fetchAndEncodeOutputImages(data).then((out) => {
+        const compactOutputs = convertToIOutputs(out, {});
+        setFullOutputs(compactOutputs);
+      });
+    }, [id, data, fullOutputs]);
 
-  if (error) {
-    return <div className="text-red-500">Error rendering output: {error.message}</div>;
-  }
+    if (error) {
+      return <div className="text-red-500">Error rendering output: {error.message}</div>;
+    }
+
+    if (!inCrossRef && exec?.ready) {
+      return (
+        <div>
+          {!fullOutputs && <div className="p-2.5">Loading...</div>}
+          {fullOutputs && <ActiveOutputRenderer id={id} data={fullOutputs} />}
+        </div>
+      );
+    }
 
-  if (!inCrossRef && registry && exec?.cell) {
     return (
-      <div ref={registry?.register(id)} data-thebe-active-ref="true">
+      <div>
         {!fullOutputs && <div className="p-2.5">Loading...</div>}
-        {fullOutputs && <ActiveOutputRenderer id={id} data={fullOutputs} />}
+        {fullOutputs && core && (
+          <PassiveOutputRenderer
+            id={id}
+            data={fullOutputs}
+            core={core}
+            kind={exec?.kind ?? SourceFileKind.Notebook}
+          ></PassiveOutputRenderer>
+        )}
       </div>
     );
-  }
-
-  return (
-    <>
-      {!fullOutputs && <div className="p-2.5">Loading...</div>}
-      {fullOutputs && core && (
-        <MemoPassiveOutputRenderer
-          id={id}
-          data={fullOutputs}
-          core={core}
-          kind={exec?.kind ?? SourceFileKind.Notebook}
-        ></MemoPassiveOutputRenderer>
-      )}
-    </>
-  );
-};
+  },
+);
diff --git a/packages/jupyter/src/output.tsx b/packages/jupyter/src/output.tsx
@@ -4,7 +4,8 @@ import type { MinifiedMimeOutput, MinifiedOutput } from 'nbtx';
 import classNames from 'classnames';
 import { SafeOutputs } from './safe';
 import { JupyterOutputs } from './jupyter';
-import { useNotebookCellExecution } from './providers';
+import { useReadyToExecute } from './providers';
+import { useMemo, useRef } from 'react';
 
 export const DIRECT_OUTPUT_TYPES = new Set(['stream', 'error']);
 
@@ -46,12 +47,15 @@ function JupyterOutput({
   nodeType?: string;
   align?: 'left' | 'center' | 'right';
 }) {
-  const exec = useNotebookCellExecution(nodeKey);
+  const ready = useReadyToExecute();
   const outputs: MinifiedOutput[] = data;
-  const allSafe = allOutputsAreSafe(outputs, DIRECT_OUTPUT_TYPES, DIRECT_MIME_TYPES);
+  const allSafe = useMemo(
+    () => allOutputsAreSafe(outputs, DIRECT_OUTPUT_TYPES, DIRECT_MIME_TYPES),
+    [outputs],
+  );
 
   let component;
-  if (allSafe && !exec?.ready) {
+  if (allSafe && !ready) {
     component = <SafeOutputs keyStub={nodeKey} outputs={outputs} />;
   } else {
     component = <JupyterOutputs id={nodeKey} outputs={outputs} />;
@@ -62,7 +66,7 @@ function JupyterOutput({
       id={identifier || undefined}
       data-mdast-node-type={nodeType}
       data-mdast-node-id={nodeKey}
-      className={classNames('max-w-full overflow-auto m-0 group not-prose relative', {
+      className={classNames('max-w-full overflow-visible m-0 group not-prose relative', {
         'text-left': !align || align === 'left',
         'text-center': align === 'center',
         'text-right': align === 'right',

diff --git a/packages/jupyter/src/providers.tsx b/packages/jupyter/src/providers.tsx
@@ -36,12 +36,6 @@ export function useComputeOptions() {
       githubBadgeUrl,
       binderBadgeUrl,
     );
-    console.debug('myst-theme:useComputeOptions:thebe', {
-      thebeFrontmatter,
-      githubBadgeUrl,
-      binderBadgeUrl,
-      thebeOptions,
-    });
     return {
       canCompute: thebeFrontmatter !== undefined && thebeFrontmatter !== false,
       thebe: thebeOptions,
@@ -122,8 +116,6 @@ export function notebookFromMdast(
   return notebook;
 }
 
-// registry[cellId]
-type CellRefRegistry = Record<string, HTMLDivElement>;
 type IdKeyMap = Record<string, string>;
 
 interface NotebookContextType {
@@ -141,9 +133,7 @@ interface NotebookContextType {
     options?: NotebookExecuteOptions | undefined,
   ) => Promise<(IThebeCellExecuteReturn | null)[]>;
   notebook: ThebeNotebook | undefined;
-  registry: CellRefRegistry;
   idkMap: IdKeyMap;
-  register: (id: string) => (el: HTMLDivElement) => void;
   restart: () => Promise<void>;
   clear: () => void;
 }
@@ -176,12 +166,10 @@ export function NotebookProvider({
     setNotebook,
   } = useNotebookBase();
 
-  const registry = useRef<CellRefRegistry>({});
   const idkMap = useRef<IdKeyMap>({});
 
   useEffect(() => {
     if (!core || !config) return;
-    registry.current = {};
     idkMap.current = {};
     if (page.kind === SourceFileKind.Notebook) {
       const nb = notebookFromMdast(
@@ -198,19 +186,6 @@ export function NotebookProvider({
     }
   }, [core, config, page]);
 
-  function register(id: string) {
-    return (el: HTMLDivElement) => {
-      if (el != null && registry.current[idkMap.current[id]] !== el) {
-        if (!el.isConnected) {
-          console.debug(`skipping ref for cell ${id} as host is not connected`);
-        } else {
-          console.debug(`new ref for cell ${id} registered`);
-          registry.current[idkMap.current[id]] = el;
-        }
-      }
-    };
-  }
-
   return (
     <NotebookContext.Provider
       value={{
@@ -223,9 +198,7 @@ export function NotebookProvider({
         executeAll,
         executeSome,
         notebook,
-        registry: registry.current,
         idkMap: idkMap.current,
-        register,
         restart: () => session?.restart() ?? Promise.resolve(),
         clear,
       }}
@@ -240,22 +213,6 @@ export function useHasNotebookProvider() {
   return notebookState !== undefined;
 }
 
-export function useCellRefRegistry() {
-  const notebookState = useContext(NotebookContext);
-  if (notebookState === undefined) return undefined;
-  return { register: notebookState.register };
-}
-
-export function useCellRef(id: string) {
-  const notebookState = useContext(NotebookContext);
-  if (notebookState === undefined) return undefined;
-
-  const { registry, idkMap } = notebookState;
-  const entry = Object.entries(notebookState.registry).find(([cellId]) => cellId === idkMap[id]);
-  console.debug('useCellRef', { id, registry, idkMap, entry });
-  return { el: entry?.[1] ?? null };
-}
-
 export function useMDASTNotebook() {
   const notebookState = useContext(NotebookContext);
   return notebookState;
@@ -271,6 +228,11 @@ export function useNotebookExecution() {
   return { ready, attached, executing, executed, errors, execute: executeAll, notebook, clear };
 }
 
+export function useReadyToExecute() {
+  const notebookState = useContext(NotebookContext);
+  return notebookState?.ready ?? false;
+}
+
 export function useNotebookCellExecution(id: string) {
   // setup a cell only executing state
   const [executing, setExecuting] = useState(false);
@@ -295,14 +257,16 @@ export function useNotebookCellExecution(id: string) {
     return execReturn;
   }
   const cell = notebook?.getCellById(cellId);
-  return {
-    kind,
-    ready,
-    cell,
-    executing,
-    notebookIsExecuting,
-    execute,
-    clear: () => cell?.clear(),
-    notebook,
-  };
+  return notebook
+    ? {
+        kind,
+        ready,
+        cell,
+        executing,
+        notebookIsExecuting,
+        execute,
+        clear: () => cell?.clear(),
+        notebook,
+      }
+    : undefined;
 }
diff --git a/packages/site/src/components/EnableCompute.tsx b/packages/site/src/components/EnableCompute.tsx
@@ -54,7 +54,7 @@ export function EnableCompute({
   }, [shutdown, navigation]);
 
   return (
-    <div className="flex mx-1 items-center">
+    <div className="flex mx-1 items-center mb-2">
       <button className={classes} onClick={() => setEnabling(true)} disabled={enabling || enabled}>
         <PowerIcon className="h-6 w-6 mx-1 inline-block align-top" title="enable compute" />
       </button>